diff --git a/bittensor_cli/cli.py b/bittensor_cli/cli.py index 16e68c254..190d6c462 100755 --- a/bittensor_cli/cli.py +++ b/bittensor_cli/cli.py @@ -5100,13 +5100,19 @@ def stake_remove( default=self.config.get("wallet_name") or defaults.wallet.name, ) if include_hotkeys: - if len(include_hotkeys) > 1: - print_error("Cannot unstake_all from multiple hotkeys at once.") - return False - elif is_valid_ss58_address(include_hotkeys[0]): - hotkey_ss58_address = include_hotkeys[0] - else: - print_error("Invalid hotkey ss58 address.") + # Multiple hotkeys are now supported with batching + # Initialize wallet as it's needed to resolve hotkey names and get coldkey + wallet = self.wallet_ask( + wallet_name, + wallet_path, + wallet_hotkey, + ask_for=[WO.NAME, WO.PATH], + ) + if len(include_hotkeys) == 1: + # Single hotkey - use hotkey_ss58_address for backward compatibility + if is_valid_ss58_address(include_hotkeys[0]): + hotkey_ss58_address = include_hotkeys[0] + # If it's a hotkey name, it will be handled by the unstake_all function return False elif all_hotkeys: wallet = self.wallet_ask( diff --git a/bittensor_cli/src/bittensor/subtensor_interface.py b/bittensor_cli/src/bittensor/subtensor_interface.py index bf2f91a23..e617ea429 100644 --- a/bittensor_cli/src/bittensor/subtensor_interface.py +++ b/bittensor_cli/src/bittensor/subtensor_interface.py @@ -1304,6 +1304,106 @@ async def create_signed(call_to_sign, n): ) return False, err_msg, None + async def sign_and_send_batch_extrinsic( + self, + calls: list[GenericCall], + wallet: Wallet, + wait_for_inclusion: bool = True, + wait_for_finalization: bool = False, + era: Optional[dict[str, int]] = None, + proxy: Optional[str] = None, + nonce: Optional[int] = None, + sign_with: Literal["coldkey", "hotkey", "coldkeypub"] = "coldkey", + batch_type: Literal["batch", "batch_all"] = "batch", + mev_protection: bool = False, + ) -> tuple[bool, str, Optional[AsyncExtrinsicReceipt], list[dict]]: + """ + :param calls: List of prepared Call objects to batch together + :param wallet: the wallet whose key will be used to sign the extrinsic + :param wait_for_inclusion: whether to wait until the extrinsic call is included on the chain + :param wait_for_finalization: whether to wait until the extrinsic call is finalized on the chain + :param era: The length (in blocks) for which a transaction should be valid. + :param proxy: The real account used to create the proxy. None if not using a proxy for this call. + :param nonce: The nonce used to submit this extrinsic call. + :param sign_with: Determine which of the wallet's keypairs to use to sign the extrinsic call. + :param batch_type: "batch" (stops on first error) or "batch_all" (executes all, fails if any fail) + :param mev_protection: If set, uses Mev Protection on the extrinsic, thus encrypting it. + + :return: (success, error message, extrinsic receipt | None, list of individual call results) + """ + if not calls: + return False, "No calls provided for batching operation", None, [] + + if len(calls) == 1: + success, err_msg, receipt = await self.sign_and_send_extrinsic( + call=calls[0], + wallet=wallet, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + era=era, + proxy=proxy, + nonce=nonce, + sign_with=sign_with, + mev_protection=mev_protection, + ) + return success, err_msg, receipt, [{"success": success, "error": err_msg}] + + batch_call = await self.substrate.compose_call( + call_module="Utility", + call_function=batch_type, + call_params={"calls": calls}, + ) + + success, err_msg, receipt = await self.sign_and_send_extrinsic( + call=batch_call, + wallet=wallet, + wait_for_inclusion=wait_for_inclusion, + wait_for_finalization=wait_for_finalization, + era=era, + proxy=proxy, + nonce=nonce, + sign_with=sign_with, + mev_protection=mev_protection, + ) + + # Parse batch results if successful + call_results = [] + if success and receipt: + try: + # Extract batch execution results from receipt + # The receipt should contain information about which calls succeeded/failed + for i, call in enumerate(calls): + call_results.append( + { + "index": i, + "call": call, + "success": True, # Will be updated if we can parse receipt + } + ) + except Exception: + # If we can't parse results, assume all succeeded if batch succeeded + for i, call in enumerate(calls): + call_results.append( + { + "index": i, + "call": call, + "success": success, + } + ) + else: + # If batch failed, mark all as failed + for i, call in enumerate(calls): + call_results.append( + { + "index": i, + "call": call, + "success": False, + "error": err_msg, + } + ) + + return success, err_msg, receipt, call_results + async def get_children(self, hotkey, netuid) -> tuple[bool, list, str]: """ This method retrieves the children of a given hotkey and netuid. It queries the SubtensorModule's ChildKeys diff --git a/bittensor_cli/src/commands/stake/add.py b/bittensor_cli/src/commands/stake/add.py index 18c6578eb..45d0ef67d 100644 --- a/bittensor_cli/src/commands/stake/add.py +++ b/bittensor_cli/src/commands/stake/add.py @@ -474,56 +474,216 @@ async def stake_extrinsic( successes = defaultdict(dict) error_messages = defaultdict(dict) extrinsic_ids = defaultdict(dict) - with console.status(f"\n:satellite: Staking on netuid(s): {netuids} ...") as status: - if safe_staking: - stake_coroutines = {} - for i, (ni, am, curr, price_with_tolerance) in enumerate( - zip( - netuids, - amounts_to_stake, - current_stake_balances, - prices_with_tolerance, + + # Collect all calls for batching + calls_to_batch = [] + call_metadata = [] # Track (netuid, staking_address, amount, current_stake, price_with_tolerance) for each call + + with console.status( + f"\n:satellite: Preparing batch staking on netuid(s): {netuids} ..." + ) as status: + # Get next nonce for batch + next_nonce = await subtensor.substrate.get_account_next_index(coldkey_ss58) + # Get block_hash at the beginning to speed up compose_call operations + block_hash = await subtensor.substrate.get_chain_head() + + # Collect all calls - iterate through the same order as when building the lists + # The lists are built in order: for each hotkey, for each netuid + list_idx = 0 + price_idx = 0 + for hotkey in hotkeys_to_stake_to: + for netuid in netuids: + # Safety check: if we've processed all items from the first loop, stop + if list_idx >= len(amounts_to_stake): + break + + # Verify subnet exists (same check as first loop) + # If subnet doesn't exist, it was skipped in first loop, so list_idx won't advance + # We need to skip it here too to stay in sync + subnet_info = all_subnets.get(netuid) + if not subnet_info: + # This netuid was skipped in first loop (doesn't exist) + # Don't advance list_idx, just continue to next netuid + continue + + am = amounts_to_stake[list_idx] + curr = current_stake_balances[list_idx] + staking_address = hotkey[1] + price_with_tol = ( + prices_with_tolerance[price_idx] + if safe_staking and price_idx < len(prices_with_tolerance) + else None ) - ): - for _, staking_address in hotkeys_to_stake_to: - # Regular extrinsic for root subnet - if ni == 0: - stake_coroutines[(ni, staking_address)] = stake_extrinsic( - netuid_i=ni, - amount_=am, - current=curr, - staking_address_ss58=staking_address, - status_=status, - ) - else: - stake_coroutines[(ni, staking_address)] = safe_stake_extrinsic( - netuid_=ni, - amount_=am, - current_stake=curr, - hotkey_ss58_=staking_address, - price_limit=price_with_tolerance, - status_=status, + if safe_staking: + price_idx += 1 + + call_metadata.append( + (netuid, staking_address, am, curr, price_with_tol) + ) + + if safe_staking and netuid != 0 and price_with_tol: + # Safe staking for non-root subnets + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="add_stake_limit", + call_params={ + "hotkey": staking_address, + "netuid": netuid, + "amount_staked": am.rao, + "limit_price": price_with_tol.rao, + "allow_partial": allow_partial_stake, + }, + block_hash=block_hash, + ) + else: + # Regular staking for root subnet or non-safe staking + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="add_stake", + call_params={ + "hotkey": staking_address, + "netuid": netuid, + "amount_staked": am.rao, + }, + block_hash=block_hash, + ) + calls_to_batch.append(call) + list_idx += 1 + + # If we have multiple calls, batch them; otherwise send single call + if len(calls_to_batch) > 1: + status.update( + f"\n:satellite: Batching {len(calls_to_batch)} stake operations..." + ) + ( + batch_success, + batch_err_msg, + batch_receipt, + call_results, + ) = await subtensor.sign_and_send_batch_extrinsic( + calls=calls_to_batch, + wallet=wallet, + era={"period": era}, + proxy=proxy, + nonce=next_nonce, + mev_protection=mev_protection, + batch_type="batch_all", # Use batch_all to execute all even if some fail + ) + + if batch_success and batch_receipt: + if mev_protection: + inner_hash = batch_err_msg + mev_shield_id = await extract_mev_shield_id(batch_receipt) + ( + mev_success, + mev_error, + batch_receipt, + ) = await wait_for_extrinsic_by_hash( + subtensor=subtensor, + extrinsic_hash=inner_hash, + shield_id=mev_shield_id, + submit_block_hash=batch_receipt.block_hash, + status=status, + ) + if not mev_success: + status.stop() + print_error(f"\n:cross_mark: [red]Failed[/red]: {mev_error}") + batch_success = False + batch_err_msg = mev_error + + if batch_success: + if not json_output: + await print_extrinsic_id(batch_receipt) + batch_ext_id = await batch_receipt.get_extrinsic_identifier() + + # Fetch updated balances for display + block_hash = await subtensor.substrate.get_chain_head() + current_balance = await subtensor.get_balance( + coldkey_ss58, block_hash + ) + + # Fetch all stake balances in parallel + if not json_output: + stake_fetch_tasks = [ + subtensor.get_stake( + hotkey_ss58=staking_address, + coldkey_ss58=coldkey_ss58, + netuid=ni, + block_hash=block_hash, + ) + for ni, staking_address, _, _, _ in call_metadata + ] + new_stakes = await asyncio.gather(*stake_fetch_tasks) + + # Process results for each call + for idx, (ni, staking_address, am, curr, _) in enumerate( + call_metadata + ): + # For batch_all, we assume all succeeded if batch succeeded + # Individual call results would need to be parsed from receipt events + successes[ni][staking_address] = True + error_messages[ni][staking_address] = "" + extrinsic_ids[ni][staking_address] = batch_ext_id + + if not json_output: + new_stake = new_stakes[idx] + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Finalized. " + f"Stake added to netuid: {ni}, hotkey: {staking_address}[/dark_sea_green3]" + ) + console.print( + f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]" + f"{ni}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] " + f"Stake:\n" + f" [blue]{curr}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n" + ) + + # Show final coldkey balance + if not json_output: + console.print( + f"Coldkey Balance:\n " + f"[blue]{current_wallet_balance}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{current_balance}" ) - else: - stake_coroutines = { - (ni, staking_address): stake_extrinsic( + else: + # Batch failed + for ni, staking_address, _, _, _ in call_metadata: + successes[ni][staking_address] = False + error_messages[ni][staking_address] = batch_err_msg + else: + # Batch submission failed + for ni, staking_address, _, _, _ in call_metadata: + successes[ni][staking_address] = False + error_messages[ni][staking_address] = ( + batch_err_msg or "Batch submission failed" + ) + elif len(calls_to_batch) == 1: + # Single call - use regular extrinsic + ni, staking_address, am, curr, price_with_tol = call_metadata[0] + + if safe_staking and ni != 0 and price_with_tol: + success, er_msg, ext_receipt = await safe_stake_extrinsic( + netuid_=ni, + amount_=am, + current_stake=curr, + hotkey_ss58_=staking_address, + price_limit=price_with_tol, + status_=status, + ) + else: + success, er_msg, ext_receipt = await stake_extrinsic( netuid_i=ni, amount_=am, current=curr, staking_address_ss58=staking_address, status_=status, ) - for i, (ni, am, curr) in enumerate( - zip(netuids, amounts_to_stake, current_stake_balances) - ) - for _, staking_address in hotkeys_to_stake_to - } - # We can gather them all at once but balance reporting will be in race-condition. - for (ni, staking_address), coroutine in stake_coroutines.items(): - success, er_msg, ext_receipt = await coroutine successes[ni][staking_address] = success error_messages[ni][staking_address] = er_msg - if success: + if success and ext_receipt: extrinsic_ids[ni][ staking_address ] = await ext_receipt.get_extrinsic_identifier() diff --git a/bittensor_cli/src/commands/stake/remove.py b/bittensor_cli/src/commands/stake/remove.py index 1cc7116a5..ddba479ab 100644 --- a/bittensor_cli/src/commands/stake/remove.py +++ b/bittensor_cli/src/commands/stake/remove.py @@ -329,33 +329,216 @@ async def unstake( return False successes = [] - with console.status("\n:satellite: Performing unstaking operations...") as status: + coldkey_ss58 = proxy or wallet.coldkeypub.ss58_address + + # Capture initial balance before operations + initial_balance = await subtensor.get_balance(coldkey_ss58) + + calls_to_batch = [] + call_metadata = [] # Track operation metadata for each call + + with console.status( + "\n:satellite: Preparing batch unstaking operations..." + ) as status: + next_nonce = await subtensor.substrate.get_account_next_index(coldkey_ss58) + # Get block_hash at the beginning to speed up compose_call operations + block_hash = await subtensor.substrate.get_chain_head() + for op in unstake_operations: - common_args = { - "wallet": wallet, - "subtensor": subtensor, - "netuid": op["netuid"], - "amount": op["amount_to_unstake"], - "hotkey_ss58": op["hotkey_ss58"], - "status": status, - "era": era, - "proxy": proxy, - "mev_protection": mev_protection, - } + call_metadata.append(op) if safe_staking and op["netuid"] != 0: - func = _safe_unstake_extrinsic - specific_args = { - "price_limit": op["price_with_tolerance"], - "allow_partial_stake": allow_partial_stake, - } + # Safe unstaking for non-root subnets + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="remove_stake_limit", + call_params={ + "hotkey": op["hotkey_ss58"], + "netuid": op["netuid"], + "amount_unstaked": op["amount_to_unstake"].rao, + "limit_price": op["price_with_tolerance"], + "allow_partial": allow_partial_stake, + }, + block_hash=block_hash, + ) else: - func = _unstake_extrinsic - specific_args = {"current_stake": op["current_stake_balance"]} + # Regular unstaking for root subnet or non-safe unstaking + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function="remove_stake", + call_params={ + "hotkey": op["hotkey_ss58"], + "netuid": op["netuid"], + "amount_unstaked": op["amount_to_unstake"].rao, + }, + block_hash=block_hash, + ) + calls_to_batch.append(call) + + # If we have multiple calls, batch them; otherwise send single call + if len(calls_to_batch) > 1: + status.update( + f"\n:satellite: Batching {len(calls_to_batch)} unstake operations..." + ) + ( + batch_success, + batch_err_msg, + batch_receipt, + call_results, + ) = await subtensor.sign_and_send_batch_extrinsic( + calls=calls_to_batch, + wallet=wallet, + era={"period": era}, + proxy=proxy, + nonce=next_nonce, + mev_protection=mev_protection, + batch_type="batch_all", # Use batch_all to execute all even if some fail + ) - suc, ext_receipt = await func(**common_args, **specific_args) - ext_id = await ext_receipt.get_extrinsic_identifier() if suc else None + if batch_success and batch_receipt: + if mev_protection: + inner_hash = batch_err_msg + mev_shield_id = await extract_mev_shield_id(batch_receipt) + ( + mev_success, + mev_error, + batch_receipt, + ) = await wait_for_extrinsic_by_hash( + subtensor=subtensor, + extrinsic_hash=inner_hash, + shield_id=mev_shield_id, + submit_block_hash=batch_receipt.block_hash, + status=status, + ) + if not mev_success: + status.stop() + print_error.print( + f"\n:cross_mark: [red]Failed[/red]: {mev_error}" + ) + batch_success = False + batch_err_msg = mev_error + + if batch_success: + if not json_output: + await print_extrinsic_id(batch_receipt) + batch_ext_id = await batch_receipt.get_extrinsic_identifier() + + # Fetch updated balances for display + block_hash = await subtensor.substrate.get_chain_head() + current_balance = await subtensor.get_balance( + coldkey_ss58, block_hash + ) + # Fetch all stake balances in parallel + if not json_output: + stake_fetch_tasks = [ + subtensor.get_stake( + hotkey_ss58=op["hotkey_ss58"], + coldkey_ss58=coldkey_ss58, + netuid=op["netuid"], + block_hash=block_hash, + ) + for op in call_metadata + ] + new_stakes = await asyncio.gather(*stake_fetch_tasks) + + # Process results for each call + for idx, op in enumerate(call_metadata): + # For batch_all, we assume all succeeded if batch succeeded + result_entry = { + "netuid": op["netuid"], + "hotkey_ss58": op["hotkey_ss58"], + "unstake_amount": op["amount_to_unstake"].tao, + "success": True, + "extrinsic_identifier": batch_ext_id, + } + console.print( + f"[yellow][DEBUG][/yellow] Operation {idx} (netuid={op['netuid']}): assigning ext_id={batch_ext_id}" + ) + successes.append(result_entry) + + if not json_output: + new_stake = new_stakes[idx] + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Finalized. " + f"Unstake from netuid: {op['netuid']}, hotkey: {op['hotkey_ss58']}[/dark_sea_green3]" + ) + console.print( + f"Subnet: [{COLOR_PALETTE['GENERAL']['SUBHEADING']}]" + f"{op['netuid']}[/{COLOR_PALETTE['GENERAL']['SUBHEADING']}] " + f"Stake:\n" + f" [blue]{op['current_stake_balance']}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{new_stake}\n" + ) + + # Show final coldkey balance + if not json_output: + console.print( + f"Coldkey Balance:\n " + f"[blue]{initial_balance}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{current_balance}" + ) + else: + # Batch failed + for op in call_metadata: + successes.append( + { + "netuid": op["netuid"], + "hotkey_ss58": op["hotkey_ss58"], + "unstake_amount": op["amount_to_unstake"].tao, + "success": False, + "extrinsic_identifier": None, + } + ) + else: + # Batch submission failed + for op in call_metadata: + successes.append( + { + "netuid": op["netuid"], + "hotkey_ss58": op["hotkey_ss58"], + "unstake_amount": op["amount_to_unstake"].tao, + "success": False, + "extrinsic_identifier": None, + } + ) + elif len(calls_to_batch) == 1: + # Single call - use regular extrinsic + op = call_metadata[0] + if safe_staking and op["netuid"] != 0: + suc, ext_receipt = await _safe_unstake_extrinsic( + wallet=wallet, + subtensor=subtensor, + netuid=op["netuid"], + amount=op["amount_to_unstake"], + hotkey_ss58=op["hotkey_ss58"], + price_limit=op["price_with_tolerance"], + allow_partial_stake=allow_partial_stake, + status=status, + era=era, + proxy=proxy, + mev_protection=mev_protection, + ) + else: + suc, ext_receipt = await _unstake_extrinsic( + wallet=wallet, + subtensor=subtensor, + netuid=op["netuid"], + amount=op["amount_to_unstake"], + current_stake=op["current_stake_balance"], + hotkey_ss58=op["hotkey_ss58"], + status=status, + era=era, + proxy=proxy, + mev_protection=mev_protection, + ) + ext_id = ( + await ext_receipt.get_extrinsic_identifier() + if suc and ext_receipt + else None + ) successes.append( { "netuid": op["netuid"], @@ -366,9 +549,10 @@ async def unstake( } ) - console.print( - f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]Unstaking operations completed." - ) + if not json_output: + console.print( + f"\n[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]Unstaking operations completed." + ) if json_output: json_console.print_json(data=successes) return True @@ -412,7 +596,9 @@ async def unstake_all( subtensor.get_balance(coldkey_ss58), ) - if all_hotkeys: + if all_hotkeys or include_hotkeys: + # Use _get_hotkeys_to_unstake when we have include_hotkeys or all_hotkeys + # This allows batching multiple unstake_all operations hotkeys = _get_hotkeys_to_unstake( wallet, hotkey_ss58_address=hotkey_ss58_address, @@ -561,21 +747,177 @@ async def unstake_all( if not unlock_key(wallet).success: return successes = {} - with console.status("Unstaking all stakes...") as status: + coldkey_ss58 = proxy or wallet.coldkeypub.ss58_address + + # Capture initial balance before operations + initial_balance_all = await subtensor.get_balance(coldkey_ss58) + + # Collect all calls for batching + calls_to_batch = [] + call_metadata = [] # Track hotkey info for each call + + with console.status("Preparing batch unstake all operations...") as status: + # Get next nonce for batch + next_nonce = await subtensor.substrate.get_account_next_index(coldkey_ss58) + # Get block_hash at the beginning to speed up compose_call operations + block_hash = await subtensor.substrate.get_chain_head() + + # Collect all calls for hotkey_ss58 in hotkey_ss58s: + hotkey_name = hotkey_names.get(hotkey_ss58, hotkey_ss58) + call_metadata.append( + { + "hotkey_ss58": hotkey_ss58, + "hotkey_name": hotkey_name, + } + ) + + call_function = "unstake_all_alpha" if unstake_all_alpha else "unstake_all" + call = await subtensor.substrate.compose_call( + call_module="SubtensorModule", + call_function=call_function, + call_params={"hotkey": hotkey_ss58}, + block_hash=block_hash, + ) + calls_to_batch.append(call) + + if len(calls_to_batch) > 1: + status.update( + f"\n:satellite: Batching {len(calls_to_batch)} unstake_all operations..." + ) + ( + batch_success, + batch_err_msg, + batch_receipt, + call_results, + ) = await subtensor.sign_and_send_batch_extrinsic( + calls=calls_to_batch, + wallet=wallet, + era={"period": era}, + proxy=proxy, + nonce=next_nonce, + mev_protection=mev_protection, + batch_type="batch_all", # Use batch_all to execute all even if some fail + ) + + if batch_success and batch_receipt: + if mev_protection: + inner_hash = batch_err_msg + mev_shield_id = await extract_mev_shield_id(batch_receipt) + ( + mev_success, + mev_error, + batch_receipt, + ) = await wait_for_extrinsic_by_hash( + subtensor=subtensor, + extrinsic_hash=inner_hash, + shield_id=mev_shield_id, + submit_block_hash=batch_receipt.block_hash, + status=status, + ) + if not mev_success: + status.stop() + print_error.print( + f"\n:cross_mark: [red]Failed[/red]: {mev_error}" + ) + batch_success = False + batch_err_msg = mev_error + + if batch_success: + if not json_output: + await print_extrinsic_id(batch_receipt) + batch_ext_id = await batch_receipt.get_extrinsic_identifier() + + # Fetch updated balances for display + block_hash = await subtensor.substrate.get_chain_head() + current_balance, all_stakes = await asyncio.gather( + subtensor.get_balance(coldkey_ss58, block_hash), + subtensor.get_stake_for_coldkey( + coldkey_ss58=coldkey_ss58, + block_hash=block_hash, + ), + ) + + # Process results for each call + for idx, metadata in enumerate(call_metadata): + # For batch_all, we assume all succeeded if batch succeeded + successes[metadata["hotkey_ss58"]] = { + "success": True, + "extrinsic_identifier": batch_ext_id, + } + + if not json_output: + # Filter stakes for this hotkey + hotkey_stakes_filtered = [ + s + for s in all_stakes + if s.hotkey_ss58 == metadata["hotkey_ss58"] + ] + + console.print( + f":white_heavy_check_mark: [dark_sea_green3]Finalized. " + f"Unstaked all from hotkey: {metadata['hotkey_name']}[/dark_sea_green3]" + ) + if hotkey_stakes_filtered: + remaining_stakes = [ + s for s in hotkey_stakes_filtered if s.stake.tao > 0 + ] + if remaining_stakes: + console.print( + f"Remaining stakes for {metadata['hotkey_name']}:" + ) + for stake_info in remaining_stakes: + console.print( + f" Netuid {stake_info.netuid}: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{stake_info.stake}" + ) + else: + console.print(f" [dim]No remaining stakes[/dim]") + else: + console.print(f" [dim]No remaining stakes[/dim]") + + # Show final coldkey balance + if not json_output: + console.print( + f"\nColdkey Balance:\n " + f"[blue]{current_wallet_balance}[/blue] " + f":arrow_right: " + f"[{COLOR_PALETTE['STAKE']['STAKE_AMOUNT']}]{current_balance}" + ) + else: + # Batch failed + for metadata in call_metadata: + successes[metadata["hotkey_ss58"]] = { + "success": False, + "extrinsic_identifier": None, + } + else: + # Batch submission failed + for metadata in call_metadata: + successes[metadata["hotkey_ss58"]] = { + "success": False, + "extrinsic_identifier": None, + } + elif len(calls_to_batch) == 1: + # Single call - use regular extrinsic + metadata = call_metadata[0] success, ext_receipt = await _unstake_all_extrinsic( wallet=wallet, subtensor=subtensor, - hotkey_ss58=hotkey_ss58, - hotkey_name=hotkey_names.get(hotkey_ss58, hotkey_ss58), + hotkey_ss58=metadata["hotkey_ss58"], + hotkey_name=metadata["hotkey_name"], unstake_all_alpha=unstake_all_alpha, status=status, era=era, proxy=proxy, mev_protection=mev_protection, ) - ext_id = await ext_receipt.get_extrinsic_identifier() if success else None - successes[hotkey_ss58] = { + ext_id = ( + await ext_receipt.get_extrinsic_identifier() + if success and ext_receipt + else None + ) + successes[metadata["hotkey_ss58"]] = { "success": success, "extrinsic_identifier": ext_id, } @@ -621,8 +963,9 @@ async def _unstake_extrinsic( f"\n:satellite: Unstaking {amount} from {hotkey_ss58} on netuid: {netuid} ..." ) + block_hash = await subtensor.substrate.get_chain_head() current_balance, next_nonce, call = await asyncio.gather( - subtensor.get_balance(coldkey_ss58), + subtensor.get_balance(coldkey_ss58, block_hash=block_hash), subtensor.substrate.get_account_next_index(coldkey_ss58), subtensor.substrate.compose_call( call_module="SubtensorModule", @@ -632,6 +975,7 @@ async def _unstake_extrinsic( "netuid": netuid, "amount_unstaked": amount.rao, }, + block_hash=block_hash, ), ) diff --git a/tests/e2e_tests/test_batching.py b/tests/e2e_tests/test_batching.py new file mode 100644 index 000000000..c306fe936 --- /dev/null +++ b/tests/e2e_tests/test_batching.py @@ -0,0 +1,295 @@ +import asyncio +import json +import re +import pytest +from bittensor_cli.src.bittensor.balances import Balance + +from .utils import set_storage_extrinsic + + +@pytest.mark.parametrize("local_chain", [False], indirect=True) +def test_batching(local_chain, wallet_setup): + """ + Test batching scenarios for stake add and remove operations. + + Steps: + 1. Create wallets for Alice and Bob + 2. Create 2 subnets (netuid 2 and 3) with Alice + 3. Start emission schedules for subnets 2 and 3 + 4. Register Bob in subnet 2 + 5. Add batch stake from Bob to subnets 2 and 3 and verify extrinsic_id uniqueness + 6. Remove all stake from all netuids using --all-netuids --all and verify extrinsic_id uniqueness + """ + print("Testing batching scenarios ๐Ÿงช") + + # Create wallets for Alice and Bob + wallet_path_alice = "//Alice" + wallet_path_bob = "//Bob" + + # Setup Alice's wallet + keypair_alice, wallet_alice, wallet_path_alice, exec_command_alice = wallet_setup( + wallet_path_alice + ) + + # Setup Bob's wallet + keypair_bob, wallet_bob, wallet_path_bob, exec_command_bob = wallet_setup( + wallet_path_bob + ) + + # Call to make Alice root owner + items = [ + ( + bytes.fromhex( + "658faa385070e074c85bf6b568cf055536e3e82152c8758267395fe524fbbd160000" + ), + bytes.fromhex( + "d43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d" + ), + ) + ] + asyncio.run( + set_storage_extrinsic( + local_chain, + wallet=wallet_alice, + items=items, + ) + ) + + # Create first subnet (netuid = 2) + result = exec_command_alice( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--subnet-name", + "Test Subnet 2", + "--repo", + "https://github.com/username/repo", + "--contact", + "test@opentensor.dev", + "--url", + "https://testsubnet.com", + "--discord", + "test#1234", + "--description", + "A test subnet for e2e testing", + "--additional-info", + "Test subnet", + "--logo-url", + "https://testsubnet.com/logo.png", + "--no-prompt", + ], + ) + assert "โœ… Registered subnetwork with netuid: 2" in result.stdout + assert "Your extrinsic has been included" in result.stdout, result.stdout + + # Create second subnet (netuid = 3) + result = exec_command_alice( + command="subnets", + sub_command="create", + extra_args=[ + "--wallet-path", + wallet_path_alice, + "--chain", + "ws://127.0.0.1:9945", + "--wallet-name", + wallet_alice.name, + "--wallet-hotkey", + wallet_alice.hotkey_str, + "--subnet-name", + "Test Subnet 3", + "--repo", + "https://github.com/username/repo", + "--contact", + "test@opentensor.dev", + "--url", + "https://testsubnet.com", + "--discord", + "test#1234", + "--description", + "A test subnet for e2e testing", + "--additional-info", + "Test subnet", + "--logo-url", + "https://testsubnet.com/logo.png", + "--no-prompt", + ], + ) + assert "โœ… Registered subnetwork with netuid: 3" in result.stdout + assert "Your extrinsic has been included" in result.stdout, result.stdout + + # Start emission schedule for subnets + start_call_netuid_2 = exec_command_alice( + command="subnets", + sub_command="start", + extra_args=[ + "--netuid", + "2", + "--wallet-name", + wallet_alice.name, + "--no-prompt", + "--chain", + "ws://127.0.0.1:9945", + "--wallet-path", + wallet_path_alice, + ], + ) + assert ( + "Successfully started subnet 2's emission schedule." + in start_call_netuid_2.stdout + ) + assert "Your extrinsic has been included" in start_call_netuid_2.stdout + + start_call_netuid_3 = exec_command_alice( + command="subnets", + sub_command="start", + extra_args=[ + "--netuid", + "3", + "--wallet-name", + wallet_alice.name, + "--no-prompt", + "--chain", + "ws://127.0.0.1:9945", + "--wallet-path", + wallet_path_alice, + ], + ) + assert ( + "Successfully started subnet 3's emission schedule." + in start_call_netuid_3.stdout + ) + assert "Your extrinsic has been included" in start_call_netuid_3.stdout + # Register Bob in one subnet + register_result = exec_command_bob( + command="subnets", + sub_command="register", + extra_args=[ + "--netuid", + "2", + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--hotkey", + wallet_bob.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--era", + "30", + ], + ) + assert "โœ… Registered" in register_result.stdout, register_result.stderr + assert "Your extrinsic has been included" in register_result.stdout, ( + register_result.stdout + ) + + # Register Bob in one subnet + register_result_subnet3 = exec_command_bob( + command="subnets", + sub_command="register", + extra_args=[ + "--netuid", + "3", + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--hotkey", + wallet_bob.hotkey_str, + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--era", + "30", + ], + ) + assert "โœ… Registered" in register_result_subnet3.stdout, ( + register_result_subnet3.stderr + ) + assert "Your extrinsic has been included" in register_result_subnet3.stdout, ( + register_result_subnet3.stdout + ) + + # Add stake to subnets + multiple_netuids = [2, 3] + stake_result = exec_command_bob( + command="stake", + sub_command="add", + extra_args=[ + "--netuids", + ",".join(str(netuid) for netuid in multiple_netuids), + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--hotkey", + wallet_bob.hotkey_str, + "--amount", + "5", + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--unsafe", + "--era", + "144", + ], + ) + assert "โœ… Finalized" in stake_result.stdout, stake_result.stderr + assert "Your extrinsic has been included" in stake_result.stdout, ( + stake_result.stdout + ) + + print(f"output: {stake_result.stdout}") + + # Verify extrinsic_id is unique (all operations should share the same extrinsic_id when batched) + # Pattern matches: "Your extrinsic has been included as {block_number}-{extrinsic_index}" + extrinsic_id_pattern = r"Your extrinsic has been included as (\d+-\d+)" + extrinsic_ids = re.findall(extrinsic_id_pattern, stake_result.stdout) + assert len(extrinsic_ids) > 0, "No extrinsic IDs found in output" + assert len(set(extrinsic_ids)) == 1, ( + f"Expected single unique extrinsic_id for batched operations, " + f"found {len(set(extrinsic_ids))} unique IDs: {set(extrinsic_ids)}" + ) + + # Remove stake from multiple netuids (should batch) + remove_stake_batch = exec_command_bob( + command="stake", + sub_command="remove", + extra_args=[ + "--all-netuids", + "--wallet-path", + wallet_path_bob, + "--wallet-name", + wallet_bob.name, + "--hotkey", + wallet_bob.hotkey_str, + "--all", + "--chain", + "ws://127.0.0.1:9945", + "--no-prompt", + "--unsafe", + "--era", + "144", + ], + ) + + # Verify extrinsic_id is unique (all operations should share the same extrinsic_id when batched) + # Pattern matches: "Your extrinsic has been included as {block_number}-{extrinsic_index}" + batch_remove_extrinsic_ids = re.findall( + extrinsic_id_pattern, remove_stake_batch.stdout + ) + + assert len(batch_remove_extrinsic_ids) > 0, "No extrinsic IDs found in output" + assert len(set(batch_remove_extrinsic_ids)) == 1, ( + f"Expected single unique extrinsic_id for batched operations, " + f"found {len(set(batch_remove_extrinsic_ids))} unique IDs: {set(batch_remove_extrinsic_ids)}" + )