Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / auto_parallel / sapp_nd / memory_estimation / _func_tracer.py: 50%

135 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-07-06 05:41 +0800

1# Copyright 2025 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""func tracer module""" 

16import ast 

17import sys 

18import inspect 

19import textwrap 

20import linecache 

21 

22 

23class _FuncTracer: 

24 """Func tracer class""" 

25 

26 def __init__(self): 

27 self.ind = 0 

28 self.trace_str = "" 

29 self.code_trees = {} 

30 self.max_sub_length = 30 

31 

32 def add2trace(self, s): 

33 """log""" 

34 self.trace_str += "\t" * self.ind + s + "\n" 

35 

36 def is_constant_expr(self, s): 

37 """whether it's a leave node""" 

38 node = ast.parse(s, mode="eval").body 

39 return isinstance(node, (ast.Constant, ast.Attribute, ast.Name)) 

40 

41 def scrap_term_symbols(self, s): 

42 """get all leave nodes from source code""" 

43 res = [] 

44 tree = ast.parse(s) 

45 for node in ast.walk(tree): 

46 if isinstance(node, ast.Name): 

47 if not any(r.startswith(node.id) for r in res): 

48 res += [node.id] 

49 elif isinstance(node, ast.Constant): 

50 res += [str(node.value)] 

51 elif isinstance(node, ast.Attribute): 

52 # Reconstruct full attribute chain like obj.attr1.attr2 

53 value = node 

54 attr_chain = [] 

55 while isinstance(value, ast.Attribute): 

56 attr_chain.insert(0, value.attr) 

57 value = value.value 

58 if isinstance(value, ast.Name): 

59 attr_chain.insert(0, value.id) 

60 res += [".".join(attr_chain)] 

61 res = sorted(res, key=len, reverse=True) 

62 return res 

63 

64 def substitute(self, s, local): 

65 """substitute leave nodes by their numerical values""" 

66 nodes = self.scrap_term_symbols(s) 

67 for n in nodes: 

68 attr_chain = n.split(".") 

69 if attr_chain[0] in local: 

70 if len(attr_chain) == 1: 

71 target = str(local[n]) 

72 if len(target) < self.max_sub_length: 

73 s = s.replace(n, target) 

74 else: 

75 s = s.replace(n, type(local[n]).__name__) 

76 else: 

77 attr = local[attr_chain[0]] 

78 for a in attr_chain[1:]: 

79 attr = getattr(attr, a) 

80 target = str(attr) 

81 if len(target) < self.max_sub_length: 

82 s = s.replace(n, target) 

83 else: 

84 s = s.replace(n, type(attr).__name__) 

85 return s 

86 

87 def fetch_node_from_lineno(self, lineno, co): 

88 """Get AST node from a line num in source code""" 

89 for node in ast.walk(self.code_trees[co]): 

90 if hasattr(node, "lineno") and node.lineno == lineno: 

91 return node 

92 return None 

93 

94 def line_tracer(self, frame, event, _): 

95 """Trace executed code line""" 

96 co = frame.f_code 

97 func_name = co.co_name 

98 

99 # Handling function call 

100 if event == "call": 

101 self.extract_ast(co) 

102 _, _, _, values = inspect.getargvalues(frame) 

103 args_str = ",".join( 

104 f"{k}={v if isinstance(v, (int, float)) else type(v).__name__}" 

105 for k, v in values.items() 

106 ) 

107 self.add2trace(f"::{func_name}({args_str})") 

108 self.ind += 1 

109 

110 # Handling instruction 

111 elif event == "line": 

112 line = linecache.getline(co.co_filename, frame.f_lineno).strip() 

113 ignored_statements = ("if ", "elif ", "else ", "for ", "while ") 

114 op_equals = ["*=", "+=", "-=", "/=", "%=", "="] 

115 # Only capture assignments and returns 

116 if not line.startswith(ignored_statements): 

117 base_str = f"{func_name}->" 

118 if "=" in line: 

119 try: 

120 ast.literal_eval(line) 

121 except (ValueError, SyntaxError): 

122 node = self.fetch_node_from_lineno(frame.f_lineno, co) 

123 if not node: 

124 return self.line_tracer 

125 line = ast.unparse(node) 

126 sign = next(op for op in op_equals if op in line) 

127 left, right = line.split(sign, 1) 

128 left, right = left.strip(), right.strip() 

129 if sign != "=": 

130 right = f"{left} {sign[0]} {right}" 

131 self.add2trace(f"{base_str} {left} = {right}") 

132 if not self.is_constant_expr(right): 

133 val = eval( # pylint: disable=eval-used 

134 right, frame.f_globals, frame.f_locals 

135 ) 

136 sub_right = self.substitute(right, frame.f_locals) 

137 base_str = f"{base_str} {left}" 

138 if not self.is_constant_expr(sub_right): 

139 self.add2trace( 

140 f"{len(base_str)*' '} = {sub_right}" 

141 ) 

142 self.add2trace(f"{len(base_str)*' '} = {val}") 

143 else: 

144 self.add2trace(f"({right}) = {val}") 

145 # Handle return statements 

146 elif line.startswith("return"): 

147 lreturn = len("return") 

148 expr = line[lreturn:].strip() 

149 try: 

150 ast.literal_eval(expr) 

151 except (ValueError, SyntaxError): 

152 node = self.fetch_node_from_lineno(frame.f_lineno, co) 

153 if not node: 

154 return self.line_tracer 

155 expr = ast.unparse(node)[lreturn:].strip() 

156 val = eval( # pylint: disable=eval-used 

157 expr, frame.f_globals, frame.f_locals 

158 ) 

159 sub_expr = self.substitute(expr, frame.f_locals) 

160 self.add2trace(f"{base_str} return {expr}") 

161 if not self.is_constant_expr(sub_expr): 

162 self.add2trace(f"{len(base_str)*' '} = {sub_expr}") 

163 self.add2trace(f"{len(base_str)*' '} = {val}") 

164 else: 

165 self.add2trace(f"{len(base_str)*' '} = {val}") 

166 

167 elif event == "return": 

168 self.ind -= 1 

169 

170 return self.line_tracer 

171 

172 def extract_ast(self, fun): 

173 """Get AST from function source code""" 

174 if fun not in self.code_trees: 

175 self.code_trees[fun] = ast.parse( 

176 textwrap.dedent(inspect.getsource(fun)) 

177 ) 

178 _, start_line = inspect.getsourcelines(fun) 

179 # Extract file's real line numbers for AST 

180 for node in ast.walk(self.code_trees[fun]): 

181 if hasattr(node, "lineno"): 

182 node.lineno = start_line + node.lineno - 1 

183 node.end_lineno = start_line + node.end_lineno - 1 

184 

185 def wrap(self, fun): 

186 """Wrapper""" 

187 

188 def tracked_fun(*args, **kwargs): 

189 """Wrapper""" 

190 sys.settrace(self.line_tracer) 

191 try: 

192 print("Tracing...") 

193 res = fun(*args, **kwargs) 

194 print(self.trace_str) 

195 self.trace_str = "" 

196 return res 

197 finally: 

198 sys.settrace(None) 

199 return res 

200 

201 return tracked_fun