Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add parallelism when performing multiple tasks at the same that can benefit from it #48

Merged
merged 16 commits into from
Jun 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions staking_deposit/cli/exit_transaction_keystore.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,8 @@ def exit_transaction_keystore(
)

folder = os.path.join(output_folder, DEFAULT_EXIT_TRANSACTION_FOLDER_NAME)
if not os.path.exists(folder):
os.mkdir(folder)

click.echo(load_text(['msg_exit_transaction_creation']))
saved_folder = export_exit_transaction_json(folder=folder, signed_exit=signed_exit)
Expand Down
102 changes: 68 additions & 34 deletions staking_deposit/cli/exit_transaction_mnemonic.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
import click
import os
import concurrent.futures

from typing import Any, Sequence
from typing import Any, Sequence, Dict
from staking_deposit.cli.existing_mnemonic import load_mnemonic_arguments_decorator
from staking_deposit.credentials import Credential
from staking_deposit.exceptions import ValidationError
Expand All @@ -23,6 +24,22 @@
from staking_deposit.utils.validation import validate_int_range, validate_validator_indices, verify_signed_exit_json


def _credential_builder(kwargs: Dict[str, Any]) -> Credential:
return Credential(**kwargs)


def _exit_exporter(kwargs: Dict[str, Any]) -> str:
credential: Credential = kwargs.pop('credential')
return credential.save_exit_transaction(**kwargs)


def _exit_verifier(kwargs: Dict[str, Any]) -> bool:
credential: Credential = kwargs.pop('credential')
kwargs['pubkey'] = credential.signing_pk.hex()
kwargs['chain_settings'] = credential.chain_setting
return verify_signed_exit_json(**kwargs)


FUNC_NAME = 'exit_transaction_mnemonic'


Expand Down Expand Up @@ -94,39 +111,56 @@ def exit_transaction_mnemonic(
key_indices = range(validator_start_index, validator_start_index + num_keys)

# We are not using CredentialList because from_mnemonic assumes key generation flow
credentials = [
Credential(
mnemonic=mnemonic,
mnemonic_password=mnemonic_password,
index=key_index,
amount=0, # Unneeded for this purpose
chain_setting=chain_settings,
hex_eth1_withdrawal_address=None
) for key_index in key_indices
]

with click.progressbar(zip(credentials, validator_indices),
label=load_text(['msg_exit_transaction_creation']),
show_percent=False,
length=num_keys,
show_pos=True) as items:
transaction_filefolders = [
credential.save_exit_transaction(validator_index=validator_index, epoch=epoch, folder=folder)
for credential, validator_index in items
]

with click.progressbar(zip(transaction_filefolders, credentials),
label=load_text(['msg_verify_exit_transaction']),
show_percent=False,
length=num_keys,
show_pos=True) as items:
if not all(
verify_signed_exit_json(file_folder=file,
pubkey=credential.signing_pk.hex(),
chain_settings=credential.chain_setting)
for file, credential in items
):
raise ValidationError(load_text(['err_verify_exit_transactions']))
credentials = []
with click.progressbar(length=num_keys, label=load_text(['msg_key_creation']),
show_percent=False, show_pos=True) as bar:

executor_kwargs = [{
'mnemonic': mnemonic,
'mnemonic_password': mnemonic_password,
'index': index,
'amount': 0,
'chain_setting': chain_settings,
'hex_eth1_withdrawal_address': None,
} for index in key_indices]

with concurrent.futures.ProcessPoolExecutor() as executor:
for credential in executor.map(_credential_builder, executor_kwargs):
credentials.append(credential)
bar.update(1)

if not os.path.exists(folder):
os.mkdir(folder)

transaction_filefolders = []
with click.progressbar(length=num_keys, label=load_text(['msg_exit_transaction_creation']),
show_percent=False, show_pos=True) as bar:

executor_kwargs = [{
'credential': credential,
'validator_index': validator_index,
'epoch': epoch,
'folder': folder,
} for credential, validator_index in zip(credentials, validator_indices)]

with concurrent.futures.ProcessPoolExecutor() as executor:
for filefolder in executor.map(_exit_exporter, executor_kwargs):
transaction_filefolders.append(filefolder)
bar.update(1)

with click.progressbar(length=num_keys, label=load_text(['msg_verify_exit_transaction']),
show_percent=False, show_pos=True) as bar:

executor_kwargs = [{
'file_folder': file,
'credential': credential,
} for file, credential in zip(transaction_filefolders, credentials)]

with concurrent.futures.ProcessPoolExecutor() as executor:
for valid_exit in executor.map(_exit_verifier, executor_kwargs):
bar.update(1)
if not valid_exit:
raise ValidationError(load_text(['err_verify_exit_transactions']))

click.echo(load_text(['msg_creation_success']) + folder)
click.pause(load_text(['msg_pause']))
34 changes: 28 additions & 6 deletions staking_deposit/cli/generate_bls_to_execution_change.py
Original file line number Diff line number Diff line change
@@ -1,15 +1,19 @@
import os
import click
import json
import concurrent.futures
from typing import (
Any,
Sequence,
Dict,
Optional
)

from eth_typing import HexAddress

from staking_deposit.credentials import (
CredentialList,
Credential
)
from staking_deposit.utils.validation import (
validate_bls_withdrawal_credentials_list,
Expand Down Expand Up @@ -48,6 +52,17 @@ def get_password(text: str) -> str:
return click.prompt(text, hide_input=True, show_default=False, type=str)


def _validate_credentials_match(kwargs: Dict[str, Any]) -> Optional[ValidationError]:
credential: Credential = kwargs.pop('credential')
bls_withdrawal_credentials: bytes = kwargs.pop('bls_withdrawal_credentials')

try:
validate_bls_withdrawal_credentials_matching(bls_withdrawal_credentials, credential)
except ValidationError as e:
return e
return None


FUNC_NAME = 'generate_bls_to_execution_change'


Expand Down Expand Up @@ -177,12 +192,19 @@ def generate_bls_to_execution_change(
)

# Check if the given old bls_withdrawal_credentials is as same as the mnemonic generated
for i, credential in enumerate(credentials.credentials):
try:
validate_bls_withdrawal_credentials_matching(bls_withdrawal_credentials_list[i], credential)
except ValidationError as e:
click.echo('\n[Error] ' + str(e))
return
with click.progressbar(length=len(credentials.credentials), label=load_text(['msg_credentials_verification']),
show_percent=False, show_pos=True) as bar:
executor_kwargs = [{
'credential': credential,
'bls_withdrawal_credentials': bls_withdrawal_credentials_list[i],
} for i, credential in enumerate(credentials.credentials)]

with concurrent.futures.ProcessPoolExecutor() as executor:
for e in executor.map(_validate_credentials_match, executor_kwargs):
bar.update(1)
if e is not None:
click.echo('\n\n[Error] ' + str(e))
return

btec_file = credentials.export_bls_to_execution_change_json(bls_to_execution_changes_folder, validator_indices)

Expand Down
115 changes: 94 additions & 21 deletions staking_deposit/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from enum import Enum
import time
import json
import concurrent.futures
from typing import Dict, List, Optional, Any, Sequence

from eth_typing import Address, HexAddress
Expand Down Expand Up @@ -230,6 +231,29 @@ def save_exit_transaction(self, validator_index: int, epoch: int, folder: str) -
return export_exit_transaction_json(folder=folder, signed_exit=signed_voluntary_exit)


def _credential_builder(kwargs: Dict[str, Any]) -> Credential:
return Credential(**kwargs)


def _keystore_exporter(kwargs: Dict[str, Any]) -> str:
credential: Credential = kwargs.pop('credential')
return credential.save_signing_keystore(**kwargs)


def _deposit_data_builder(credential: Credential) -> Dict[str, bytes]:
return credential.deposit_datum_dict


def _keystore_verifier(kwargs: Dict[str, Any]) -> bool:
credential: Credential = kwargs.pop('credential')
return credential.verify_keystore(**kwargs)


def _bls_to_execution_change_builder(kwargs: Dict[str, Any]) -> Dict[str, bytes]:
credential: Credential = kwargs.pop('credential')
return credential.get_bls_to_execution_change_dict(**kwargs)


class CredentialList:
"""
A collection of multiple Credentials, one for each validator.
Expand All @@ -253,23 +277,52 @@ def from_mnemonic(cls,
f"The number of keys ({num_keys}) doesn't equal to the corresponding deposit amounts ({len(amounts)})."
)
key_indices = range(start_index, start_index + num_keys)
with click.progressbar(key_indices, label=load_text(['msg_key_creation']),
show_percent=False, show_pos=True) as indices:
return cls([Credential(mnemonic=mnemonic, mnemonic_password=mnemonic_password,
index=index, amount=amounts[index - start_index], chain_setting=chain_setting,
hex_eth1_withdrawal_address=hex_eth1_withdrawal_address,
use_pbkdf2=use_pbkdf2)
for index in indices])

credentials: List[Credential] = []
with click.progressbar(length=num_keys, label=load_text(['msg_key_creation']),
show_percent=False, show_pos=True) as bar:
executor_kwargs = [{
'mnemonic': mnemonic,
'mnemonic_password': mnemonic_password,
'index': index,
'amount': amounts[index - start_index],
'chain_setting': chain_setting,
'hex_eth1_withdrawal_address': hex_eth1_withdrawal_address,
'use_pbkdf2': use_pbkdf2,
} for index in key_indices]

with concurrent.futures.ProcessPoolExecutor() as executor:
for credential in executor.map(_credential_builder, executor_kwargs):
credentials.append(credential)
bar.update(1)
return cls(credentials)

def export_keystores(self, password: str, folder: str) -> List[str]:
with click.progressbar(self.credentials, label=load_text(['msg_keystore_creation']),
show_percent=False, show_pos=True) as credentials:
return [credential.save_signing_keystore(password=password, folder=folder) for credential in credentials]
filefolders: List[str] = []
with click.progressbar(length=len(self.credentials), label=load_text(['msg_keystore_creation']),
show_percent=False, show_pos=True) as bar:
executor_kwargs = [{
'credential': credential,
'password': password,
'folder': folder,
} for credential in self.credentials]

with concurrent.futures.ProcessPoolExecutor() as executor:
for filefolder in executor.map(_keystore_exporter, executor_kwargs):
filefolders.append(filefolder)
bar.update(1)
return filefolders

def export_deposit_data_json(self, folder: str) -> str:
with click.progressbar(self.credentials, label=load_text(['msg_depositdata_creation']),
show_percent=False, show_pos=True) as credentials:
deposit_data = [cred.deposit_datum_dict for cred in credentials]
deposit_data = []
with click.progressbar(length=len(self.credentials), label=load_text(['msg_depositdata_creation']),
show_percent=False, show_pos=True) as bar:

with concurrent.futures.ProcessPoolExecutor() as executor:
for datum_dict in executor.map(_deposit_data_builder, self.credentials):
deposit_data.append(datum_dict)
bar.update(1)

filefolder = os.path.join(folder, 'deposit_data-%i.json' % time.time())
with open(filefolder, 'w') as f:
json.dump(deposit_data, f, default=lambda x: x.hex())
Expand All @@ -278,17 +331,37 @@ def export_deposit_data_json(self, folder: str) -> str:
return filefolder

def verify_keystores(self, keystore_filefolders: List[str], password: str) -> bool:
with click.progressbar(zip(self.credentials, keystore_filefolders),
all_valid_keystores = True
with click.progressbar(length=len(self.credentials),
label=load_text(['msg_keystore_verification']),
length=len(self.credentials), show_percent=False, show_pos=True) as items:
return all(credential.verify_keystore(keystore_filefolder=filefolder, password=password)
for credential, filefolder in items)
show_percent=False, show_pos=True) as bar:
executor_kwargs = [{
'credential': credential,
'keystore_filefolder': fileholder,
'password': password,
} for credential, fileholder in zip(self.credentials, keystore_filefolders)]

with concurrent.futures.ProcessPoolExecutor() as executor:
for valid_keystore in executor.map(_keystore_verifier, executor_kwargs):
all_valid_keystores &= valid_keystore
bar.update(1)

return all_valid_keystores

def export_bls_to_execution_change_json(self, folder: str, validator_indices: Sequence[int]) -> str:
with click.progressbar(self.credentials, label=load_text(['msg_bls_to_execution_change_creation']),
show_percent=False, show_pos=True) as credentials:
bls_to_execution_changes = [cred.get_bls_to_execution_change_dict(validator_indices[i])
for i, cred in enumerate(credentials)]
bls_to_execution_changes = []
with click.progressbar(length=len(self.credentials), label=load_text(['msg_bls_to_execution_change_creation']),
show_percent=False, show_pos=True) as bar:

executor_kwargs = [{
'credential': credential,
'validator_index': validator_indices[i],
} for i, credential in enumerate(self.credentials)]

with concurrent.futures.ProcessPoolExecutor() as executor:
for bls_to_execution_change in executor.map(_bls_to_execution_change_builder, executor_kwargs):
bls_to_execution_changes.append(bls_to_execution_change)
bar.update(1)

filefolder = os.path.join(folder, 'bls_to_execution_change-%i.json' % time.time())
with open(filefolder, 'w') as f:
Expand Down
2 changes: 2 additions & 0 deletions staking_deposit/deposit.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import click
import socket
import sys
from multiprocessing import freeze_support

from staking_deposit.cli.existing_mnemonic import existing_mnemonic
from staking_deposit.cli.exit_transaction_keystore import exit_transaction_keystore
Expand Down Expand Up @@ -91,6 +92,7 @@ def cli(ctx: click.Context, language: str, non_interactive: bool, ignore_connect


def run() -> None:
freeze_support() # Needed when running under Windows in a frozen bundle
check_python_version()
cli()

Expand Down
3 changes: 0 additions & 3 deletions staking_deposit/exit_transaction.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,9 +48,6 @@ def export_exit_transaction_json(folder: str, signed_exit: SignedVoluntaryExit)
signed_exit_json.update({'message': message})
signed_exit_json.update({'signature': '0x' + signed_exit.signature.hex()}) # type: ignore[attr-defined]

if not os.path.exists(folder):
os.mkdir(folder)

filefolder = os.path.join(
folder,
'signed_exit_transaction-%s-%i.json' % (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
"arg_exit_transaction_mnemonic_output_folder": {
"help": "The folder path where the exit transactions will be saved to. Pointing to `./exit_transactions` by default."
},
"msg_key_creation": "Creating your keys:\t",
"msg_exit_transaction_creation": "Creating your exit transactions:\t",
"msg_verify_exit_transaction": "Verifying your exit transactions:\t",
"err_verify_exit_transactions": "\nThere was a problem verifying your exit transactions.\nPlease try again",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
"help": "The folder path for the keystore(s). Pointing to `./bls_to_execution_changes` by default."
},
"msg_key_creation": "Creating your SignedBLSToExecutionChange.",
"msg_credentials_verification": "Verifying your withdrawal credentials.",
"msg_creation_success": "\nSuccess!\nYour SignedBLSToExecutionChange JSON file can be found at: ",
"msg_pause": "\n\nPress any key.",
"err_verify_btec": "Failed to verify the bls_to_execution_change JSON files."
Expand Down
Loading