Source code for rhodent.utils.env

from __future__ import annotations


from os import getenv


[docs] def rhodent_getenv(variable) -> str: """ Get value of environment variable, or default. Possible environment variables are: - `RHODENT_RESPONSE_MAX_MEM` - Value in units of MiB. When constructing response, try to limit memory usage below this value. - `RHODENT_RESPONSE_MAX_MEM_PER_RANK` - Value in units of MiB. Ignored if `RHODENT_RESPONSE_MAX_MEM` is set. - `RHODENT_REDISTRIBUTE_MAXSIZE` - Maximal number of elements in single redistribute call. Larger arrays are split. Parameters ---------- variable Name of variable, without the ``'RHODENT_'`` prefix. Returns ------- Value of variable as string. """ defaults = {'REDISTRIBUTE_MAXSIZE': '1e7', 'RESPONSE_MAX_MEM': '', 'RESPONSE_MAX_MEM_PER_RANK': '1000', } if variable not in defaults.keys(): allowed = '\n'.join([f' RHODENT_{var}' for var in defaults.keys()]) raise ValueError(f'Environment variable RHODENT_{variable} is unknown. ' f'Allowed variables are:\n{allowed}') return getenv(f'RHODENT_{variable}', defaults[variable])
def get_float(variable) -> float: strvalue = rhodent_getenv(variable) try: value = float(strvalue) except ValueError: raise ValueError(f'Expected environment variable RHODENT_{variable} to ' f'be castable to float. Value is {strvalue!r}.') return value def get_response_max_mem(comm_size: int): if rhodent_getenv('RESPONSE_MAX_MEM') == '': return get_float('RESPONSE_MAX_MEM_PER_RANK') * comm_size return get_float('RESPONSE_MAX_MEM')