Coverage for mindsdb / utilities / profiler / profiler.py: 43%

1import time 

2from datetime import datetime, timezone 

3from functools import wraps 

4 

5import mindsdb.utilities.hooks as hooks 

6from mindsdb.utilities.config import Config 

7from mindsdb.utilities.context import context as ctx 

8 

9 

10def _get_current_node(profiling: dict) -> dict: 

11 """ return the node that the pointer points to 

12 

13 Args: 

14 profiling (dict): whole profiling data 

15 

16 Returns: 

17 dict: current node 

18 """ 

19 current_node = profiling['tree'] 

20 for child_index in profiling['pointer']: 

21 current_node = current_node['children'][child_index] 

22 return current_node 

23 

24 

25def start_node(tag: str): 

26 """ Add a new node to profiling 

27 

28 Args: 

29 tag (str): name of new node 

30 """ 

31 profiling = ctx.profiling 

32 new_node = { 

33 'start_at': time.perf_counter(), 

34 'start_at_thread': time.thread_time(), 

35 'start_at_process': time.process_time(), 

36 'stop_at': None, 

37 'name': tag, 

38 'children': [] 

39 } 

40 if profiling['pointer'] is None: 

41 if profiling['level'] != 1: 

42 # profiling was activated not in the root of nodes tree 

43 return 

44 profiling['pointer'] = [] 

45 profiling['tree'] = new_node 

46 profiling['tree']['time_start_at'] = datetime.now(timezone.utc) 

47 else: 

48 current_node = _get_current_node(profiling) 

49 profiling['pointer'].append(len(current_node['children'])) 

50 current_node['children'].append(new_node) 

51 

52 

53def stop_current_node(): 

54 """ Mark current node as completed and move pointer up 

55 """ 

56 profiling = ctx.profiling 

57 if profiling['pointer'] is None: 

58 # profiling was activated not in the root of nodes tree 

59 return 

60 current_node = _get_current_node(profiling) 

61 current_node['stop_at'] = time.perf_counter() 

62 current_node['stop_at_thread'] = time.thread_time() 

63 current_node['stop_at_process'] = time.process_time() 

64 if len(profiling['pointer']) > 0: 

65 profiling['pointer'] = profiling['pointer'][:-1] 

66 else: 

67 if ctx.profiling['enabled'] is True: 

68 _send_profiling_results() 

69 profiling['pointer'] = None 

70 

71 

72def set_meta(**kwargs): 

73 """ Add any additional info to profiling data 

74 

75 Args: 

76 **kwargs (dict): metadata to add 

77 """ 

78 if profiling_enabled() is True: 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true

79 ctx.profiling.update(kwargs) 

80 

81 

82def _send_profiling_results(): 

83 """ Send profiling results to storage 

84 """ 

85 ctx.profiling['company_id'] = ctx.company_id 

86 ctx.profiling['hostname'] = Config().get('aws_meta_data', {}).get('public-hostname', '?') 

87 ctx.profiling['instance_id'] = Config().get('aws_meta_data', {}).get('instance-id', '?') 

88 hooks.send_profiling_results(ctx.profiling) 

89 

90 

91def enable(): 

92 ctx.profiling['enabled'] = True 

93 

94 

95def disable(): 

96 ctx.profiling['enabled'] = False 

97 

98 

99def profiling_enabled(): 

100 try: 

101 return ctx.profiling['enabled'] is True 

102 except AttributeError: 

103 return False 

104 

105 

106def start(tag): 

107 """ add new node to profiling data 

108 """ 

109 ctx.profiling['level'] += 1 

110 if profiling_enabled() is True: 110 ↛ 111line 110 didn't jump to line 111 because the condition on line 110 was never true

111 start_node(tag) 

112 

113 

114def stop(): 

115 """ finalize current node and move pointer up 

116 """ 

117 ctx.profiling['level'] -= 1 

118 if profiling_enabled() is True: 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true

119 stop_current_node() 

120 

121 

122class Context: 

123 def __init__(self, tag): 

124 self.tag = tag 

125 

126 def __enter__(self): 

127 start(self.tag) 

128 

129 def __exit__(self, exc_type, exc_value, traceback): 

130 stop() 

131 

132 

133def profile(tag: str = None): 

134 def decorator(function): 

135 @wraps(function) 

136 def wrapper(*args, **kwargs): 

137 if profiling_enabled() is True: 137 ↛ 138line 137 didn't jump to line 138 because the condition on line 137 was never true

138 with Context(tag or f'{function.__name__}|{function.__module__}'): 

139 result = function(*args, **kwargs) 

140 else: 

141 result = function(*args, **kwargs) 

142 return result 

143 return wrapper 

144 return decorator