wehub-resource-sync
259 lines
8.0 KiB
Python
259 lines
8.0 KiB
Python
# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# 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.
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
from argparse import ArgumentParser
|
|
|
|
import paddle
|
|
from paddle.base.log_helper import get_logger
|
|
|
|
_logger = get_logger(
|
|
__name__, logging.INFO, fmt='%(asctime)s-%(levelname)s: %(message)s'
|
|
)
|
|
|
|
|
|
color_map = {
|
|
"forward": "thread_state_running", # RGB: 126, 200, 148
|
|
"backward": "rail_idle", # RGB: 238, 142, 0
|
|
"optimizer": "rail_response", # RGB: 238, 142, 0
|
|
"default": "thread_state_unknown", # RGB: 199, 155, 125
|
|
}
|
|
|
|
ignore_job_type = ["recv_forward", "send_backward"]
|
|
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser()
|
|
device_count = paddle.device.cuda.device_count()
|
|
all_devices = ",".join([str(i) for i in range(device_count)])
|
|
parser.add_argument("--devices", type=str, default=all_devices)
|
|
parser.add_argument("--log_dir", type=str, required=True)
|
|
parser.add_argument("--multi_machine", action="store_true")
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
def process_job_log(log_data, device_id, multi_machine_idx=-1):
|
|
log_pattern = r'.*?Profiler Info: Job \((\d+)\), type = (\w+), micro_batch_id = (\d+), job_start_time = (\d+.\d+), job_end_time = (\d+.\d+)'
|
|
matches = re.findall(log_pattern, log_data)
|
|
events = []
|
|
last_end_time = None
|
|
|
|
step_times = []
|
|
step_start_time = 0
|
|
step_end_time = 0
|
|
|
|
start_job_type = ""
|
|
|
|
for i, match in enumerate(matches):
|
|
job_id, job_type, micro_batch_id, job_start_time, job_end_time = match
|
|
|
|
if job_type in ignore_job_type:
|
|
continue
|
|
|
|
if job_type != "default" and start_job_type == "":
|
|
start_job_type = job_type
|
|
|
|
start_time = float(job_start_time.strip()) * 1000
|
|
end_time = float(job_end_time.strip()) * 1000
|
|
|
|
is_start_time_recorded = 0
|
|
|
|
if job_type == start_job_type and micro_batch_id == "0":
|
|
if step_start_time != 0:
|
|
step_times.append([step_start_time, step_end_time])
|
|
step_start_time = start_time
|
|
|
|
step_end_time = end_time
|
|
|
|
tid_name = (
|
|
"GPU" + str(device_id)
|
|
if multi_machine_idx == -1
|
|
else "GPU"
|
|
+ str(device_id)
|
|
+ "(machine:"
|
|
+ str(multi_machine_idx)
|
|
+ ")"
|
|
)
|
|
event_start = {
|
|
"name": job_type + "_" + str(job_id),
|
|
"cat": job_type,
|
|
"ph": "B",
|
|
"ts": start_time,
|
|
"pid": 0,
|
|
"tid": tid_name,
|
|
}
|
|
event_end = {
|
|
"name": job_type + "_" + str(job_id),
|
|
"cat": job_type,
|
|
"ph": "E",
|
|
"pid": 0,
|
|
"ts": end_time,
|
|
"tid": tid_name,
|
|
}
|
|
if job_type in color_map:
|
|
event_start["cname"] = color_map[job_type]
|
|
event_end["cname"] = color_map[job_type]
|
|
|
|
events.append(event_start)
|
|
events.append(event_end)
|
|
|
|
last_end_time = end_time
|
|
|
|
step_times.append([step_start_time, step_end_time])
|
|
return events, step_times
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
all_events = []
|
|
step_infos = []
|
|
start_step = 0
|
|
machine_num = 1
|
|
|
|
def process_one_machine_log(log_dir, multi_machine_idx=-1):
|
|
for device_id in args.devices.split(","):
|
|
_logger.info(f"Process device {device_id}")
|
|
device_id = int(device_id)
|
|
log_file = os.path.join(log_dir, "workerlog." + str(device_id))
|
|
with open(log_file, "r") as f:
|
|
log_data = f.read()
|
|
|
|
start_step_pattern = (
|
|
r'.*?Schedule Profiler start at step (\d+) and end at step.*'
|
|
)
|
|
start_step_match = re.findall(start_step_pattern, log_data)
|
|
start_step = (
|
|
int(start_step_match[0]) if len(start_step_match) > 0 else 0
|
|
)
|
|
|
|
events, step_times = process_job_log(
|
|
log_data, device_id, multi_machine_idx
|
|
)
|
|
all_events.extend(events)
|
|
for i, info in enumerate(step_times):
|
|
if len(step_infos) <= i:
|
|
step_infos.append([float("inf"), float("-inf")])
|
|
step_infos[i][0] = min(step_infos[i][0], info[0])
|
|
step_infos[i][1] = max(step_infos[i][1], info[1])
|
|
return start_step
|
|
|
|
if args.multi_machine:
|
|
multi_machine_dirs = os.listdir(args.log_dir)
|
|
multi_machine_dirs = [
|
|
os.path.join(args.log_dir, d)
|
|
for d in multi_machine_dirs
|
|
if d.startswith("machine")
|
|
and os.path.isdir(os.path.join(args.log_dir, d))
|
|
]
|
|
machine_num = len(multi_machine_dirs)
|
|
for i, d in enumerate(multi_machine_dirs):
|
|
_logger.info(f"Process machine {i}")
|
|
start_step = max(process_one_machine_log(d, i), start_step)
|
|
else:
|
|
start_step = process_one_machine_log(args.log_dir)
|
|
|
|
for i, info in enumerate(step_infos):
|
|
start_time = info[0]
|
|
if i > 0:
|
|
start_time = max(start_time, step_infos[i - 1][1])
|
|
event_start = {
|
|
"name": "step" + str(i + start_step),
|
|
"cat": "step",
|
|
"ph": "B",
|
|
"ts": start_time,
|
|
"pid": 0,
|
|
"tid": "Step",
|
|
"cname": color_map["default"],
|
|
}
|
|
event_end = {
|
|
"name": "step" + str(i + start_step),
|
|
"cat": "step",
|
|
"ph": "E",
|
|
"ts": info[1],
|
|
"pid": 0,
|
|
"tid": "Step",
|
|
"cname": color_map["default"],
|
|
}
|
|
|
|
all_events.append(event_start)
|
|
all_events.append(event_end)
|
|
|
|
save_path = os.path.join(args.log_dir, "pipeline_profile.json")
|
|
with open(save_path, "w") as f:
|
|
f.write(json.dumps({"traceEvents": all_events}))
|
|
_logger.info(f"Save pipeline profile to {save_path}")
|
|
|
|
# support Perfetto format
|
|
save_path = os.path.join(args.log_dir, "pipeline_profile_perfetto.json")
|
|
all_events.extend(
|
|
[
|
|
{
|
|
"args": {"name": "STEP"},
|
|
"cat": "__metadata",
|
|
"name": "thread_name",
|
|
"ph": "M",
|
|
"pid": 0,
|
|
"tid": 2333,
|
|
"ts": 0,
|
|
}
|
|
]
|
|
)
|
|
|
|
for i in range(machine_num):
|
|
for j in range(len(args.devices.split(","))):
|
|
if machine_num > 1:
|
|
name = f"GPU:{j}(machine:{i})"
|
|
tid = i * len(args.devices.split(",")) + j + 2334
|
|
else:
|
|
name = f"GPU:{j}"
|
|
tid = j + 2334
|
|
all_events.extend(
|
|
[
|
|
{
|
|
"args": {"name": name},
|
|
"cat": "__metadata",
|
|
"name": "thread_name",
|
|
"ph": "M",
|
|
"pid": 0,
|
|
"tid": tid,
|
|
"ts": 0,
|
|
}
|
|
]
|
|
)
|
|
|
|
json_str = json.dumps({"traceEvents": all_events})
|
|
json_str = json_str.replace('"Step"', '2333')
|
|
|
|
for i in range(machine_num):
|
|
for j in range(len(args.devices.split(","))):
|
|
if machine_num > 1:
|
|
json_str = json_str.replace(
|
|
f'"GPU{j}(machine:{i})"',
|
|
f'{i * len(args.devices.split(",")) + j + 2334}',
|
|
)
|
|
else:
|
|
json_str = json_str.replace(f'"GPU{j}"', f'{j + 2334}')
|
|
|
|
with open(save_path, "w") as f:
|
|
f.write(json_str)
|
|
_logger.info(f"Save pipeline profile to {save_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|