from loguru import logger as log
import psutil
import numpy as np
[docs]
def get_budget(nr_sources, nr_rows, nr_chans, nr_corrs, data_type, cb_args,
fudge_factor=1.25, row2source_ratio=100):
systmem = float(psutil.virtual_memory()[0])
if not cb_args.num_workers:
nrthreads = psutil.cpu_count()
else:
nrthreads = cb_args.num_workers
log.info('-' * 50)
log.info('Budgeting')
log.info('-' * 50)
log.info('system RAM = {0:.2f} GB', systmem / 1024**3)
log.info('nr of logical CPUs = {0:d}', nrthreads)
log.info('nr sources = {0:d}', nr_sources)
log.info('nr rows = {0:d}', nr_rows)
log.info('nr chans = {0:d}', nr_chans)
log.info('nr corrs = {0:d}', nr_corrs)
data_type = {'complex': 'complex64', 'dcomplex': 'complex128'}[data_type]
data_bytes = np.dtype(data_type).itemsize
bytes_per_row = nr_chans * nr_corrs * data_bytes
memory_per_row = bytes_per_row * fudge_factor
if cb_args.model_chunks and cb_args.row_chunks:
rows_per_chunk = cb_args.row_chunks
sources_per_chunk = cb_args.model_chunks
strat_type = "(user settings)"
elif not cb_args.model_chunks and not cb_args.row_chunks:
if cb_args.memory_fraction > 1 or cb_args.memory_fraction <= 0:
raise ValueError('The memory fraction must be a number in the '
'interval (0,1]. You have set it to {0:f} '
'.'.format(cb_args.memory_fraction))
allowed_rows_per_thread = (systmem * cb_args.memory_fraction /
(memory_per_row * nrthreads))
rows_per_chunk = int(min(nr_rows, allowed_rows_per_thread))
sources_per_chunk = int(min(nr_sources,
rows_per_chunk / row2source_ratio))
strat_type = "(auto settings)"
else:
raise ValueError('For now you must set both row and source chunk, or '
'leave both unset (=0); '
'you cannot set only one of them.')
log.info('sources per chunk = {0:.0f} {1}', sources_per_chunk, strat_type)
log.info('rows per chunk = {0:.0f} {1}', rows_per_chunk, strat_type)
memory_usage = (rows_per_chunk * memory_per_row * nrthreads
+ sources_per_chunk)
log.info('expected memory usage = {0:.2f} GB', memory_usage / 1024**3)
return rows_per_chunk, sources_per_chunk