150 lines
5.2 KiB
Python
150 lines
5.2 KiB
Python
# coding=utf-8
|
|
from __future__ import unicode_literals, absolute_import
|
|
import json
|
|
from argparse import ArgumentParser
|
|
from collections import defaultdict
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument('req_json', type=str,help="reorder request file")
|
|
parser.add_argument('res_json', type=str,help="reorder result file")
|
|
parser.add_argument('trace_dst_file',type=str,help="write chrome trace file path")
|
|
|
|
args = parser.parse_args()
|
|
with open(args.req_json) as fin:
|
|
req_json = json.load(fin)
|
|
with open(args.res_json) as fin:
|
|
res_json = json.load(fin)
|
|
uuid2node ={}
|
|
nodes = req_json["nodes"]
|
|
for node in nodes:
|
|
uuid2node[node["uuid"]] = node
|
|
# cat,dur,ts,
|
|
edge_map = {}
|
|
edges = req_json["edges"]
|
|
for edge in edges:
|
|
from_id = edge["from"]
|
|
to_id = edge["to"]
|
|
if uuid2node[from_id]["typename"]!="compute" and \
|
|
uuid2node[to_id]["typename"]!="compute": # comm start node
|
|
edge_map[from_id] = edge
|
|
traceEvents = []
|
|
chrome_trace_json_format = {
|
|
"distributedInfo":{
|
|
"rank":0
|
|
},
|
|
|
|
}
|
|
stackFrames = {}
|
|
# antopt format
|
|
if "finalized_node_start_end_time" in res_json:
|
|
finalized_node_start_end_time = res_json["finalized_node_start_end_time"]
|
|
# prepare comm edge
|
|
for node_start_end in finalized_node_start_end_time:
|
|
node_uuid = node_start_end["node_uuid"]
|
|
if "_" in node_uuid:#TODO: edge,
|
|
continue
|
|
node = uuid2node[node_uuid]
|
|
if node["typename"]=="compute":
|
|
traceEvents.append({
|
|
"id":node_uuid,
|
|
"cat": "kernel",
|
|
"dur": node_start_end["end_time"]-node_start_end["start_time"],
|
|
"ts": node_start_end["start_time"],
|
|
"pid": 0,
|
|
"tid": 0,
|
|
"name": node["name"],
|
|
"ph": "X",
|
|
"args": {
|
|
"cost":node["cost"],
|
|
"uuid":node_uuid
|
|
|
|
}
|
|
})
|
|
else:
|
|
name = "ncclKernel_"+node["name"]
|
|
if node_uuid not in edge_map:
|
|
continue
|
|
cost = edge_map[node_uuid]["cost"]
|
|
traceEvents.append({
|
|
"id":node_uuid,
|
|
"cat": "kernel",
|
|
"dur": cost,
|
|
"ts": node_start_end["start_time"],
|
|
"pid": 0,
|
|
"tid": 1,
|
|
"name": name,
|
|
"ph": "X",
|
|
"args": {
|
|
"cost": cost,
|
|
"uuid":node_uuid
|
|
|
|
}
|
|
})
|
|
elif "nodes" in res_json:
|
|
# ortools format
|
|
# prepare communication edge
|
|
node2deps = defaultdict(list)
|
|
for edge in req_json["edges"]:
|
|
node2deps[edge["to"]].append(edge["from"])
|
|
for node in res_json["nodes"]:
|
|
node_before_solve = uuid2node[node["uuid"]]
|
|
# overwrite
|
|
node["opcode"] = node_before_solve["opcode"]
|
|
uuid2node[node["uuid"]]=node
|
|
def normalize_dt(time):
|
|
return time
|
|
for node in res_json["nodes"]:
|
|
uuid = node["uuid"]
|
|
deps = node2deps[uuid]
|
|
|
|
if node["typename"]=="compute":
|
|
traceEvents.append({
|
|
"id": uuid,
|
|
"cat": "kernel",
|
|
"dur": normalize_dt(node["endTime"]-node["startTime"]),
|
|
"ts": normalize_dt(node["startTime"]),
|
|
"pid": 0,
|
|
"tid": 0,
|
|
"name": node["name"],
|
|
"ph":"X",
|
|
"args":{
|
|
"cost": normalize_dt(node["endTime"]-node["startTime"]),
|
|
"uuid":uuid,
|
|
"opcode":node["opcode"],
|
|
"obj":{
|
|
"id_ref":node2deps[uuid]
|
|
},
|
|
},
|
|
})
|
|
else:
|
|
name = "ncclKernel_"+node["name"]
|
|
if uuid in edge_map:
|
|
cost = edge_map[uuid]["cost"]
|
|
else:
|
|
# compute link to communication
|
|
continue
|
|
traceEvents.append({
|
|
"id":uuid,
|
|
"cat": "kernel",
|
|
"dur":normalize_dt(cost),
|
|
"ts": normalize_dt(node["startTime"]),
|
|
"pid": 0,
|
|
"tid": 1,
|
|
"name": name,
|
|
"ph": "X",
|
|
"args": {
|
|
"cost": normalize_dt(cost),
|
|
"uuid":uuid,
|
|
"opcode":node["opcode"]
|
|
|
|
}
|
|
})
|
|
count = 0
|
|
|
|
chrome_trace_json_format["traceEvents"] = traceEvents
|
|
chrome_trace_json_format["stackFrames"] = stackFrames
|
|
with open(args.trace_dst_file, "w") as fout:
|
|
json.dump(chrome_trace_json_format, fout,indent=2)
|
|
|
|
if __name__ == '__main__':
|
|
main() |