Skip to content

Commit

Permalink
Add pool concurrency config
Browse files Browse the repository at this point in the history
  • Loading branch information
evgeny-stakewise committed Sep 11, 2023
1 parent 669b3fe commit 9d8b2e9
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 7 deletions.
21 changes: 17 additions & 4 deletions src/commands/create_keys.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,13 +46,21 @@
type=str,
callback=validate_eth_address,
)
@click.option(
'--concurrency',
help='Number of processes in a pool.',
# do not prompt
type=int,
)
@click.command(help='Creates the validator keys from the mnemonic.')
# pylint: disable-next=too-many-arguments
def create_keys(
mnemonic: str,
count: int,
vault: HexAddress,
data_dir: str,
per_keystore_password: bool,
concurrency: int | None,
) -> None:
config = VaultConfig(vault, Path(data_dir))
config.load(mnemonic)
Expand All @@ -67,16 +75,18 @@ def create_keys(
mnemonic=mnemonic,
count=count,
start_index=config.mnemonic_next_index,
concurrency=concurrency,
)
deposit_data = _export_deposit_data_json(
credentials=credentials, filename=str(deposit_data_file)
credentials=credentials, filename=str(deposit_data_file), concurrency=concurrency
)

_export_keystores(
credentials=credentials,
keystores_dir=str(keystores_dir),
password_file=str(password_file),
per_keystore_password=per_keystore_password,
concurrency=concurrency,
)

config.increment_mnemonic_index(count)
Expand All @@ -88,13 +98,15 @@ def create_keys(
)


def _export_deposit_data_json(credentials: list[Credential], filename: str) -> str:
def _export_deposit_data_json(
credentials: list[Credential], filename: str, concurrency: int | None = None
) -> str:
with click.progressbar(
length=len(credentials),
label='Generating deposit data JSON\t\t',
show_percent=False,
show_pos=True,
) as progress_bar, Pool() as pool:
) as progress_bar, Pool(processes=concurrency) as pool:
results = [
pool.apply_async(
cred.deposit_datum_dict,
Expand All @@ -117,6 +129,7 @@ def _export_keystores(
keystores_dir: str,
password_file: str,
per_keystore_password: bool,
concurrency: int | None = None,
) -> None:
makedirs(path.abspath(keystores_dir), exist_ok=True)
if not per_keystore_password:
Expand All @@ -126,7 +139,7 @@ def _export_keystores(
label='Exporting validator keystores\t\t',
show_percent=False,
show_pos=True,
) as progress_bar, Pool() as pool:
) as progress_bar, Pool(processes=concurrency) as pool:
results = [
pool.apply_async(
cred.save_signing_keystore,
Expand Down
10 changes: 8 additions & 2 deletions src/common/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -115,16 +115,22 @@ def deposit_datum_dict(self) -> dict[str, bytes]:

class CredentialManager:
@staticmethod
# pylint: disable-next=too-many-arguments
def generate_credentials(
network: str, vault: HexAddress, mnemonic: str, count: int, start_index: int
network: str,
vault: HexAddress,
mnemonic: str,
count: int,
start_index: int,
concurrency: int | None = None,
) -> list[Credential]:
credentials: list[Credential] = []
with click.progressbar(
length=count,
label='Creating validator keys:\t\t',
show_percent=False,
show_pos=True,
) as progress_bar, Pool() as pool:
) as progress_bar, Pool(processes=concurrency) as pool:

def bar_updated(result):
progress_bar.update(len(result))
Expand Down
4 changes: 4 additions & 0 deletions src/config/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ class Settings(metaclass=Singleton):
ipfs_fetch_endpoints: list[str]
validators_fetch_chunk_size: int
sentry_dsn: str
load_keystores_concurrency: int | None

# pylint: disable-next=too-many-arguments,too-many-locals
def set(
Expand Down Expand Up @@ -121,6 +122,9 @@ def set(
'VALIDATORS_FETCH_CHUNK_SIZE', default=100, cast=int
)
self.sentry_dsn = decouple_config('SENTRY_DSN', default='')
self.load_keystores_concurrency = decouple_config(
'LOAD_KEYSTORES_CONCURRENCY', default=None, cast=lambda x: int(x) if x else None
)

@property
def network_config(self) -> NetworkConfig:
Expand Down
2 changes: 1 addition & 1 deletion src/validators/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ def load_keystores() -> Keystores | None:

keystore_files = list_keystore_files()
logger.info('Loading keystores from %s...', settings.keystores_dir)
with Pool() as pool:
with Pool(processes=settings.load_keystores_concurrency) as pool:
# pylint: disable-next=unused-argument
def _stop_pool(*args, **kwargs):
pool.close()
Expand Down

0 comments on commit 9d8b2e9

Please sign in to comment.