Coverage for /opt/cloud/slavespace/usr1/096471637100f3de0fcfc01072822a80/ut/src/asys/collect/graph/graph_collect.py: 97%

1#!/usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3# ---------------------------------------------------------------------------- 

4# Copyright (c) 2025 Huawei Technologies Co., Ltd. 

5# 

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

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

8# You may obtain a copy of the License at 

9# 

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

11# 

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

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

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

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

16# limitations under the License. 

17# ---------------------------------------------------------------------------- 

18 

19import os 

20 

21from params import ParamDict 

22from common import consts 

23from common import log_debug, log_warning 

24from common import FileOperate as f 

25from common.file_operate import COPY_MODE, MOVE_MODE 

26from drv import EnvVarName 

27 

28__all__ = ["collect_graph"] 

29 

30 

31def collect_graph_files(source_dir, target_dir): 

32 ge_res = [] 

33 tf_res = [] 

34 for root, _, files in os.walk(source_dir): 

35 for file in files: 

36 if file.startswith("ge_onnx") and file.endswith(".pbtxt"): 

37 ge_res.append(os.path.join(root, file)) 

38 elif file.startswith("ge_proto") and file.endswith(".txt"): 

39 ge_res.append(os.path.join(root, file)) 

40 elif file.startswith("TF_GeOp") and file.endswith(".pbtxt"): 

41 tf_res.append(os.path.join(root, file)) 

42 

43 if not (ge_res or tf_res): 

44 return False 

45 

46 ge_target_dir_path = os.path.join(target_dir, "ge") 

47 tf_target_dir_path = os.path.join(target_dir, "tf") 

48 

49 ret = True 

50 for file_path in ge_res: 

51 ge_file_ret = f.collect_file_to_dir(file_path, ge_target_dir_path, COPY_MODE) 

52 ret = ret and ge_file_ret 

53 

54 for file_path in tf_res: 

55 tf_file_ret = f.collect_file_to_dir(file_path, tf_target_dir_path, COPY_MODE) 

56 ret = ret and tf_file_ret 

57 return ret 

58 

59 

60def collect_cmd_graph_files(graph_target_dir): 

61 collect_graph_path_list = [] 

62 

63 task_dir_path = ParamDict().get_arg("task_dir") # task_dir checked in set_arg 

64 if task_dir_path and f.check_dir(task_dir_path): 

65 collect_graph_path_list.append(task_dir_path) 

66 

67 env_var = EnvVarName() 

68 if env_var.npu_collect_path: 

69 log_debug("Get env NPU_COLLECT_PATH successfully, add NPU_COLLECT_PATH to graph collect path.") 

70 collect_graph_path_list.append(env_var.npu_collect_path) 

71 if env_var.dump_graph_path: 

72 log_debug("Get env DUMP_GRAPH_PATH successfully, add DUMP_GRAPH_PATH to graph collect path.") 

73 collect_graph_path_list.append(env_var.dump_graph_path) 

74 if env_var.work_path: 

75 log_debug("Get env ASCEND_WORK_PATH successfully, add ASCEND_WORK_PATH to graph collect path.") 

76 collect_graph_path_list.append(env_var.work_path) 

77 collect_graph_path_list.append(env_var.current_path) 

78 

79 for collect_graph_path in collect_graph_path_list: 

80 ret = collect_graph_files(collect_graph_path, graph_target_dir) 

81 if ret: 

82 return True 

83 return False 

84 

85 

86def collect_graph(output_root_path): 

87 if (ParamDict().get_command() == consts.launch_cmd) and (not ParamDict().get_ini("graph") == "1"): # 1: open 

88 log_debug("graph is not set on, not collect graph files") 

89 return 

90 

91 ret = False 

92 graph_target_dir = os.path.join(output_root_path, "dfx", "graph") 

93 if ParamDict().get_command() == consts.launch_cmd: 

94 npu_collect_path = os.path.join(ParamDict().asys_output_timestamp_dir, "npu_collect_intermediates") 

95 graph_source_dir = os.path.join(npu_collect_path, "extra-info", "graph") 

96 if f.check_dir(graph_source_dir): 

97 log_debug("Graph source check success, path={}.".format(graph_source_dir)) 

98 ret = f.collect_dir(graph_source_dir, graph_target_dir, MOVE_MODE) 

99 else: 

100 ret = collect_cmd_graph_files(graph_target_dir) 

101 

102 if not ret: 

103 log_warning("Graph collect failed.")