Coverage for rhodent/utils/env.py: 74%

1from __future__ import annotations 

2 

3 

4from os import getenv 

5 

6 

7def rhodent_getenv(variable) -> str: 

8 """ Get value of environment variable, or default. 

9 

10 Possible environment variables are: 

11 

12 - `RHODENT_RESPONSE_MAX_MEM` - Value in units of MiB. When constructing 

13 response, try to limit memory usage below this value. 

14 - `RHODENT_RESPONSE_MAX_MEM_PER_RANK` - Value in units of MiB. 

15 Ignored if `RHODENT_RESPONSE_MAX_MEM` is set. 

16 - `RHODENT_REDISTRIBUTE_MAXSIZE` - Maximal number of elements in single redistribute call. 

17 Larger arrays are split. 

18 

19 Parameters 

20 ---------- 

21 variable 

22 Name of variable, without the ``'RHODENT_'`` prefix. 

23 

24 Returns 

25 ------- 

26 Value of variable as string. 

27 """ 

28 

29 defaults = {'REDISTRIBUTE_MAXSIZE': '1e7', 

30 'RESPONSE_MAX_MEM': '', 

31 'RESPONSE_MAX_MEM_PER_RANK': '1000', 

32 } 

33 

34 if variable not in defaults.keys(): 

35 allowed = '\n'.join([f' RHODENT_{var}' for var in defaults.keys()]) 

36 raise ValueError(f'Environment variable RHODENT_{variable} is unknown. ' 

37 f'Allowed variables are:\n{allowed}') 

38 

39 return getenv(f'RHODENT_{variable}', defaults[variable]) 

40 

41 

42def get_float(variable) -> float: 

43 strvalue = rhodent_getenv(variable) 

44 try: 

45 value = float(strvalue) 

46 except ValueError: 

47 raise ValueError(f'Expected environment variable RHODENT_{variable} to ' 

48 f'be castable to float. Value is {strvalue!r}.') 

49 

50 return value 

51 

52 

53def get_response_max_mem(comm_size: int): 

54 if rhodent_getenv('RESPONSE_MAX_MEM') == '': 

55 return get_float('RESPONSE_MAX_MEM_PER_RANK') * comm_size 

56 

57 return get_float('RESPONSE_MAX_MEM')