mindspore/tests/st/debugger/dump_test_utils.py

67 lines
2.8 KiB
Python

# Copyright 2021 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""
Utils for testing offline debugger.
"""
import os
import tempfile
import bisect
import csv
import numpy as np
def build_dump_structure(path, tensor_name_list, tensor_list, net_name, tensor_info_list):
"""Build dump file structure from tensor_list."""
ranks_run_history = {}
temp_dir = tempfile.mkdtemp(prefix=net_name, dir=path)
for tensor_name, tensor, tensor_info in zip(tensor_name_list, tensor_list, tensor_info_list):
slot = str(tensor_info.slot)
iteration = str(tensor_info.iteration)
rank_id = str(tensor_info.rank_id)
root_graph_id = str(tensor_info.root_graph_id)
is_output = str(tensor_info.is_output)
graphs_run_history = ranks_run_history.get(rank_id)
if graphs_run_history is None:
graphs_run_history = {}
ranks_run_history[rank_id] = graphs_run_history
if root_graph_id not in graphs_run_history:
graphs_run_history[root_graph_id] = [iteration]
if iteration not in graphs_run_history[root_graph_id]:
bisect.insort(graphs_run_history[root_graph_id], iteration)
path = os.path.join(temp_dir, "rank_" + rank_id, net_name, root_graph_id, iteration)
os.makedirs(path, exist_ok=True)
if is_output == "True":
file_name = f'{tensor_name}.output.{slot}.DefaultFormat.npy'
else:
file_name = f'{tensor_name}.input.{slot}.DefaultFormat.npy'
full_path = os.path.join(path, file_name)
np.save(full_path, tensor)
build_global_execution_order(temp_dir, ranks_run_history)
return temp_dir
def build_global_execution_order(path, ranks_run_history):
"""Build global execution order."""
for rank_id in ranks_run_history.keys():
exec_order_path = path + "/rank_" + rank_id + "/" + "execution_order"
os.makedirs(exec_order_path, exist_ok=True)
for graph in ranks_run_history[rank_id].keys():
full_path = os.path.join(exec_order_path, "ms_global_execution_order_graph_" + graph + ".csv")
with open(full_path, 'w+', newline='') as csv_file:
write = csv.writer(csv_file)
write.writerows(ranks_run_history[rank_id][graph])