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

136 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-08-21 04:29 +0800

1# Copyright 2025-2026 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 inspect 

18import linecache 

19import sys 

20import textwrap 

21import weakref 

22 

23 

24class _FuncTracer: 

25 """Func tracer class""" 

26 

27 def __init__(self): 

28 self.ind = 0 

29 self.trace_str = "" 

30 self.code_trees = weakref.WeakKeyDictionary() 

31 self.max_sub_length = 30 

32 

33 def add2trace(self, s): 

34 """log""" 

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

36 

37 def is_constant_expr(self, s): 

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

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

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

41 

42 def scrap_term_symbols(self, s): 

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

44 res = [] 

45 tree = ast.parse(s) 

46 for node in ast.walk(tree): 

47 if isinstance(node, ast.Name): 

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

49 res += [node.id] 

50 elif isinstance(node, ast.Constant): 

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

52 elif isinstance(node, ast.Attribute): 

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

54 value = node 

55 attr_chain = [] 

56 while isinstance(value, ast.Attribute): 

57 attr_chain.insert(0, value.attr) 

58 value = value.value 

59 if isinstance(value, ast.Name): 

60 attr_chain.insert(0, value.id) 

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

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

63 return res 

64 

65 def substitute(self, s, local): 

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

67 nodes = self.scrap_term_symbols(s) 

68 for n in nodes: 

69 attr_chain = n.split(".") 

70 if attr_chain[0] in local: 

71 if len(attr_chain) == 1: 

72 target = str(local[n]) 

73 if len(target) < self.max_sub_length: 

74 s = s.replace(n, target) 

75 else: 

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

77 else: 

78 attr = local[attr_chain[0]] 

79 for a in attr_chain[1:]: 

80 attr = getattr(attr, a) 

81 target = str(attr) 

82 if len(target) < self.max_sub_length: 

83 s = s.replace(n, target) 

84 else: 

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

86 return s 

87 

88 def fetch_node_from_lineno(self, lineno, co): 

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

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

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

92 return node 

93 return None 

94 

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

96 """Trace executed code line""" 

97 co = frame.f_code 

98 func_name = co.co_name 

99 

100 # Handling function call 

101 if event == "call": 

102 self.extract_ast(co) 

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

104 args_str = ",".join( 

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

106 for k, v in values.items() 

107 ) 

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

109 self.ind += 1 

110 

111 # Handling instruction 

112 elif event == "line": 

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

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

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

116 # Only capture assignments and returns 

117 if not line.startswith(ignored_statements): 

118 base_str = f"{func_name}->" 

119 if "=" in line: 

120 try: 

121 ast.literal_eval(line) 

122 except (ValueError, SyntaxError): 

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

124 if not node: 

125 return self.line_tracer 

126 line = ast.unparse(node) 

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

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

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

130 if sign != "=": 

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

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

133 if not self.is_constant_expr(right): 

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

135 right, frame.f_globals, frame.f_locals 

136 ) 

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

138 base_str = f"{base_str} {left}" 

139 if not self.is_constant_expr(sub_right): 

140 self.add2trace( 

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

142 ) 

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

144 else: 

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

146 # Handle return statements 

147 elif line.startswith("return"): 

148 lreturn = len("return") 

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

150 try: 

151 ast.literal_eval(expr) 

152 except (ValueError, SyntaxError): 

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

154 if not node: 

155 return self.line_tracer 

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

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

158 expr, frame.f_globals, frame.f_locals 

159 ) 

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

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

162 if not self.is_constant_expr(sub_expr): 

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

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

165 else: 

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

167 

168 elif event == "return": 

169 self.ind -= 1 

170 

171 return self.line_tracer 

172 

173 def extract_ast(self, fun): 

174 """Get AST from function source code""" 

175 if fun not in self.code_trees: 

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

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

178 ) 

179 _, start_line = inspect.getsourcelines(fun) 

180 # Extract file's real line numbers for AST 

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

182 if hasattr(node, "lineno"): 

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

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

185 

186 def wrap(self, fun): 

187 """Wrapper""" 

188 

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

190 """Wrapper""" 

191 sys.settrace(self.line_tracer) 

192 try: 

193 print("Tracing...") 

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

195 print(self.trace_str) 

196 self.trace_str = "" 

197 return res 

198 finally: 

199 sys.settrace(None) 

200 return res 

201 

202 return tracked_fun