Merge pull request #48 from remyroy/multiprocess

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

staking_deposit/cli/exit_transaction_keystore.py

@@ -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)

staking_deposit/cli/exit_transaction_mnemonic.py

@@ -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
@@ -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'
 
 
@@ -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']))

staking_deposit/cli/generate_bls_to_execution_change.py

@@ -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,
@@ -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'
 
 
@@ -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)
 

staking_deposit/credentials.py

@@ -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
@@ -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.
@@ -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())
@@ -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:

staking_deposit/deposit.py

@@ -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
@@ -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()
 

staking_deposit/exit_transaction.py

@@ -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' % (

staking_deposit/intl/en/cli/exit_transaction_mnemonic.json

@@ -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",

staking_deposit/intl/en/cli/generate_bls_to_execution_change.json

@@ -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."