Coverage for gpaw/new/timer.py: 48%

1from contextlib import contextmanager 

2from functools import wraps 

3from io import StringIO 

4from typing import Callable, TypeVar, Union, overload, Optional 

5from gpaw import GPAW_TRACE 

6import inspect 

7 

8 

9class GlobalTimer: 

10 def __init__(self): 

11 from gpaw.utilities.timing import Timer 

12 self._timers = [Timer()] 

13 

14 @contextmanager 

15 def context(self, timer): 

16 # XXX We need to decide what "timer contexts" are, 

17 # at least that will be necessary if we want a default 

18 # behaviour which would then be in effect during runs of the 

19 # test suite. 

20 if not GPAW_TRACE: 

21 raise RuntimeError('You need to set environment variable ' 

22 'GPAW_TRACE=1 in order to utilize ' 

23 'tracing via GlobalTimer.') 

24 self._timers.append(timer) 

25 try: 

26 yield 

27 finally: 

28 self._timers.pop() 

29 

30 def start(self, name, **kwargs): 

31 timer = self._timers[-1] 

32 n_params = len(inspect.signature(timer.start).parameters) 

33 if n_params == 1: 

34 timer.start(name) 

35 else: 

36 timer.start(name, **kwargs) 

37 

38 def stop(self, name=None, **kwargs): 

39 timer = self._timers[-1] 

40 n_params = len(inspect.signature(timer.stop).parameters) 

41 if n_params == 1: 

42 timer.stop(name) 

43 else: 

44 timer.stop(name, **kwargs) 

45 

46 def tostring(self): 

47 buf = StringIO() 

48 self._timers[-1].write(out=buf) 

49 return buf.getvalue() 

50 

51 

52T = TypeVar('T') 

53F = Callable[..., T] 

54 

55 

56@overload 

57def _trace(meth: F, **timer_params) -> F: 

58 ... 

59 

60 

61@overload 

62def _trace(meth: None = None, **timer_params) -> Callable[[F], F]: 

63 ... 

64 

65 

66def _trace(meth: Optional[F] = None, 

67 **timer_params) -> Union[F, Callable[[F], F]]: 

68 """Decorator for telling global timer to trace a function or method.""" 

69 

70 def get_wrapper(method: Callable[..., T]) -> Callable[..., T]: 

71 modname = method.__module__ 

72 methname = method.__qualname__ 

73 name = f'{modname}.{methname}' 

74 

75 @wraps(method) 

76 def wrapper(*args, **kwargs) -> T: 

77 global_timer.start(name, **timer_params) 

78 try: 

79 return method(*args, **kwargs) 

80 finally: 

81 global_timer.stop(**timer_params) 

82 

83 return wrapper 

84 

85 if meth: 

86 return get_wrapper(meth) 

87 

88 return get_wrapper 

89 

90 

91def dummy_trace(meth: Union[Callable[..., T], None] = None, 

92 **timer_params) -> Callable[..., T]: 

93 if meth: 

94 return meth 

95 

96 def wrapper(method): 

97 return method 

98 

99 return wrapper 

100 

101 

102@contextmanager 

103def _tracectx(name, gpu=False): 

104 global_timer.start(name, gpu=gpu) 

105 try: 

106 yield 

107 finally: 

108 global_timer.stop(gpu=gpu) 

109 

110 

111@contextmanager 

112def dummy_tracectx(name, gpu=False): 

113 yield 

114 

115 

116trace = _trace if GPAW_TRACE else dummy_trace 

117tracectx = _tracectx if GPAW_TRACE else dummy_tracectx 

118global_timer = GlobalTimer()