【torch profiler】新增torch profiler解析-MindStudio-Service-Profiler-AtomGit

Mms_service_profiler/data_source/__init__.py+3-2

	@@ -14,7 +14,8 @@
14	# See the Mulan PSL v2 for more details.	14	# See the Mulan PSL v2 for more details.
15	# -------------------------------------------------------------------------	15	# -------------------------------------------------------------------------
16		16
17	-__all__ = ["DBDataSource", "MsprofDataSource", "MsptiDataSource"]	17	+__all__ = ["DBDataSource", "MsprofDataSource", "MsptiDataSource", "TorchProfilerDataSource"]
18	from .db_data_source import DBDataSource	18	from .db_data_source import DBDataSource
19	from .msprof_data_source import MsprofDataSource	19	from .msprof_data_source import MsprofDataSource
20	-from .mspti_data_source import MsptiDataSource	20	+from .mspti_data_source import MsptiDataSource
		21	+from .torch_profiler_data_source import TorchProfilerDataSource

Mms_service_profiler/data_source/msprof_data_source.py+2-0

	@@ -323,6 +323,8 @@ class MsprofDataSource(BaseDataSource):
323	self.clear_last_msprof_output(cur_path)	323	self.clear_last_msprof_output(cur_path)
324	self.run_msprof_command(command)	324	self.run_msprof_command(command)
325	filepaths = self.get_filepaths(prof_path, file_filter)	325	filepaths = self.get_filepaths(prof_path, file_filter)
		326	+ if filepaths.get("msprof", {}):
		327	+ filepaths["msprof"] = [path for path in filepaths["msprof"] if "ascend_pt" not in path]
326	try:	328	try:
327	data = self.load_prof(filepaths)	329	data = self.load_prof(filepaths)
328	except Exception as ex:	330	except Exception as ex:

Ams_service_profiler/data_source/torch_profiler_data_source.py+81-0

@@ -0,0 +1,81 @@
	1	+# -------------------------------------------------------------------------
	2	+# This file is part of the MindStudio project.
	3	+# Copyright (c) 2025 Huawei Technologies Co.,Ltd.
	4	+#
	5	+# MindStudio is licensed under Mulan PSL v2.
	6	+# You can use this software according to the terms and conditions of the Mulan PSL v2.
	7	+# You may obtain a copy of Mulan PSL v2 at:
	8	+#
	9	+# http://license.coscl.org.cn/MulanPSL2
	10	+#
	11	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	12	+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	13	+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	14	+# See the Mulan PSL v2 for more details.
	15	+# -------------------------------------------------------------------------
	16	+
	17	+import os
	18	+import json
	19	+import subprocess
	20	+from pathlib import Path
	21	+
	22	+from ms_service_profiler.data_source.base_data_source import BaseDataSource, Task
	23	+from ms_service_profiler.utils.log import logger
	24	+from ms_service_profiler.exporters.utils import find_all_file_complete
	25	+
	26	+
	27	+@Task.register("data_source:torch_profiler")
	28	+class TorchProfilerDataSource(BaseDataSource):
	29	+
	30	+ @classmethod
	31	+ def outputs(cls):
	32	+ return ["data_source:torch_profiler"]
	33	+
	34	+ @classmethod
	35	+ def get_prof_paths(cls, input_path: str):
	36	+ filepaths = []
	37	+ for dp in Path(input_path).rglob("*/_ascend_pt"):
	38	+ if dp.is_dir():
	39	+ filepaths.append(dp)
	40	+
	41	+ return filepaths
	42	+
	43	+ @classmethod
	44	+ def is_need_torchprofiler(cls, full_path):
	45	+ torch_profiler_path = os.path.join(full_path, 'ASCEND_PROFILER_OUTPUT')
	46	+ if not os.path.isdir(torch_profiler_path):
	47	+ return True
	48	+
	49	+ return False
	50	+
	51	+
	52	+ @classmethod
	53	+ def run_torch_profiler_parse(cls, full_path):
	54	+ try:
	55	+ import torch
	56	+ import torch_npu
	57	+ from torch_npu.profiler.profiler import analyse
	58	+ except ImportError as e:
	59	+ logger.warning("Required module not available: %s", str(e))
	60	+ return None
	61	+
	62	+ try:
	63	+ result = analyse(profiler_path=full_path)
	64	+ logger.info("Successfully parsed msprof data from: %s", full_path)
	65	+ return result
	66	+
	67	+ except Exception as e:
	68	+ logger.error("Failed to parse msprof data: %s", str(e))
	69	+ return None
	70	+
	71	+
	72	+ def load(self, prof_path):
	73	+ file_filter = {
	74	+ "torch_profiler": "trace_view.json",
	75	+ }
	76	+ cur_path = str(prof_path)
	77	+ if self.is_need_torchprofiler(cur_path):
	78	+ self.run_torch_profiler_parse(cur_path)
	79	+ filepaths = self.get_filepaths(prof_path, file_filter)
	80	+ filepaths['tx_data_df'] = None
	81	+ return filepaths

Mms_service_profiler/exporters/exporter_trace.py+30-14

	@@ -66,11 +66,11 @@ class ExporterTrace(TaskExporterBase):
66		66
67	@classmethod	67	@classmethod
68	def depends(cls):	68	def depends(cls):
69	- return ["pipeline:service", "pipeline:mspti"]	69	+ return ["pipeline:service", "pipeline:mspti", "data_source:torch_profiler"]
70		70
71	@classmethod	71	@classmethod
72	@timer(logger.debug)	72	@timer(logger.debug)
73	- def export(cls, data, mspti) -> None:	73	+ def export(cls, data, mspti, torch_profiler) -> None:
74	if 'db' not in cls.args.format and 'json' not in cls.args.format:	74	if 'db' not in cls.args.format and 'json' not in cls.args.format:
75	return	75	return
76		76
	@@ -129,38 +129,55 @@ class ExporterTrace(TaskExporterBase):
129	mspti_data = {"traceEvents": tarce_events_list}	129	mspti_data = {"traceEvents": tarce_events_list}
130	merged_data = merge_json_data(merged_data, [mspti_data])	130	merged_data = merge_json_data(merged_data, [mspti_data])
131		131
132	- if 'json' in cls.args.format:	132	+ if torch_profiler:
133	- save_trace_data_into_json(merged_data, output)	133	+ torch_profiler_prof_list = []
		134	+ for index, prof_path in enumerate(torch_profiler):
		135	+ torch_profiler_prof, _ = load_single_prof(prof_path, index)
		136	+ torch_profiler_prof_list.extend(torch_profiler_prof.get("traceEvents"))
134		137
135	if 'db' in cls.args.format:	138	if 'db' in cls.args.format:
136	logger.info('Start write trace data to db')	139	logger.info('Start write trace data to db')
137	create_sqlite_tables(TRACE_TABLE_DEFINITIONS)	140	create_sqlite_tables(TRACE_TABLE_DEFINITIONS)
138	save_trace_data_into_db(merged_data)	141	save_trace_data_into_db(merged_data)
139	logger.info('Write trace data to db success')	142	logger.info('Write trace data to db success')
		143	+
		144	+ if 'json' in cls.args.format:
		145	+ if torch_profiler:
		146	+ chrome_events = merged_data.get('traceEvents', [])
		147	+ merged_events = torch_profiler_prof_list + chrome_events
		148	+ merged_data = {'traceEvents': merged_events}
		149	+ save_trace_data_into_json(merged_data, output)
140		150
141	def do_export(self):	151	def do_export(self):
142	- data, mspti = self.get_depends_result("pipeline:service", None), self.get_depends_result("pipeline:mspti", None)	152	+ data, mspti, torch_profiler = (
		153	+ self.get_depends_result("pipeline:service", None),
		154	+ self.get_depends_result("pipeline:mspti", None),
		155	+ self.get_depends_result("data_source:torch_profiler", None),
		156	+ )
		157	+ torch_profiler = torch_profiler.get("torch_profiler") if torch_profiler is not None else None
143		158
144	- if self.task_index == 0 and (data is not None or mspti is not None):	159	+ has_any_data = data is not None or mspti is not None or torch_profiler is not None
145	- data_list = self.gather((None, None), dst=0)	160	+ if self.task_index == 0 and has_any_data:
		161	+ data_list = self.gather((None, None, None), dst=0)
146	else:	162	else:
147	- data_list = self.gather((data, mspti), dst=0)	163	+ data_list = self.gather((data, mspti, torch_profiler), dst=0)
148		164
149	if data_list is None:	165	if data_list is None:
150	return None	166	return None
151		167
152	- if self.task_index == 0 and (data is not None or mspti is not None):	168	+ if self.task_index == 0 and has_any_data:
153	- data_list[0] = (data, mspti)	169	+ data_list[0] = (data, mspti, torch_profiler)
154		170
155	# 使用列表推导式过滤并提取非None值	171	# 使用列表推导式过滤并提取非None值
156	all_data = [item[0] for item in data_list if item is not None and item[0] is not None]	172	all_data = [item[0] for item in data_list if item is not None and item[0] is not None]
157	all_mspti = [item[1] for item in data_list if item is not None and item[1] is not None]	173	all_mspti = [item[1] for item in data_list if item is not None and item[1] is not None]
		174	+ all_profilers = [item[2] for item in data_list if item is not None and item[2] is not None]
158		175
159	# 获取第一个非None的值	176	# 获取第一个非None的值
160	valid_data = all_data[0] if all_data else None	177	valid_data = all_data[0] if all_data else None
161	valid_mspti = all_mspti[0] if all_mspti else None	178	valid_mspti = all_mspti[0] if all_mspti else None
162		179
163	- self.export(valid_data, valid_mspti)	180	+ self.export(valid_data, valid_mspti, all_profilers)
164	return None	181	return None
165		182
166		183
	@@ -260,7 +277,7 @@ def write_trace_data_to_file(trace_data, output):
260	logger.error(f"Error raise from exporter trace, json dump failed. message: {e}")	277	logger.error(f"Error raise from exporter trace, json dump failed. message: {e}")
261		278
262	with ms_open(output, "w") as f:	279	with ms_open(output, "w") as f:
263	- f.write('{"traceEvents":[')	280	+ f.write('[')
264	for index, content2 in enumerate(results):	281	for index, content2 in enumerate(results):
265	if len(content2) < 2: # 确保至少有 2 个字符	282	if len(content2) < 2: # 确保至少有 2 个字符
266	continue	283	continue
	@@ -268,14 +285,13 @@ def write_trace_data_to_file(trace_data, output):
268	if index != len(results) - 1:	285	if index != len(results) - 1:
269	f.write(',')	286	f.write(',')
270		287
271	- f.write("]}")	288	+ f.write("]")
272		289
273	logger.info(f"Written trace data successfully. at {output}")	290	logger.info(f"Written trace data successfully. at {output}")
274		291
275		292
276	def save_trace_data_into_json(trace_data, output):	293	def save_trace_data_into_json(trace_data, output):
277	file_path = os.path.join(output, 'chrome_tracing.json')	294	file_path = os.path.join(output, 'chrome_tracing.json')
278	-
279	write_trace_data_to_file(trace_data.get("traceEvents", []), file_path)	295	write_trace_data_to_file(trace_data.get("traceEvents", []), file_path)
280		296
281		297

Mtest/ut/python/test_ms_service_profiler/test_exporter_trace.py+6-6

	@@ -205,11 +205,11 @@ def test_exporter_initialize():
205	@mock.patch('ms_service_profiler.exporters.exporter_trace.save_trace_data_into_json')	205	@mock.patch('ms_service_profiler.exporters.exporter_trace.save_trace_data_into_json')
206	def test_exporter_export(mock_save, mock_create, mock_data, mock_mspti):	206	def test_exporter_export(mock_save, mock_create, mock_data, mock_mspti):
207	# 模拟 create_trace_events 返回值	207	# 模拟 create_trace_events 返回值
208	- mock_create.return_value = {'traceEvents': []}	208	+ mock_create.return_value = {"traceEvents": []}
209	mock_save.return_value = None # 模拟保存行为	209	mock_save.return_value = None # 模拟保存行为
210		210
211	ExporterTrace.initialize(mock.Mock(output_path='/tmp', format=['json']))	211	ExporterTrace.initialize(mock.Mock(output_path='/tmp', format=['json']))
212	- ExporterTrace.export(mock_data, mock_mspti)	212	+ ExporterTrace.export(mock_data, mock_mspti, mock_data)
213		213
214	# 验证 create_trace_events 被调用一次	214	# 验证 create_trace_events 被调用一次
215	mock_create.assert_called_once()	215	mock_create.assert_called_once()
	@@ -470,8 +470,8 @@ class MockFile:
470	def test_save_trace_data_into_json(mock_logger, mock_ms_open):	470	def test_save_trace_data_into_json(mock_logger, mock_ms_open):
471	fake_file = MockFile()	471	fake_file = MockFile()
472	mock_ms_open.return_value = fake_file	472	mock_ms_open.return_value = fake_file
473	- save_trace_data_into_json({'traceEvents': [{}]}, 'output')	473	+ save_trace_data_into_json({}, 'output')
474	- assert (fake_file.read() == '{"traceEvents":[{}]}')	474	+ assert (fake_file.read() == '[]')
475	mock_logger.info.assert_called_once()	475	mock_logger.info.assert_called_once()
476		476
477		477
	@@ -482,9 +482,9 @@ def test_save_trace_data_into_json_exception(mock_logger, mock_dumps, mock_ms_op
482	mock_dumps.side_effect = Exception('Test exception')	482	mock_dumps.side_effect = Exception('Test exception')
483	fake_file = MockFile()	483	fake_file = MockFile()
484	mock_ms_open.return_value = fake_file	484	mock_ms_open.return_value = fake_file
485	- save_trace_data_into_json({'traceEvents': [123]}, 'output')	485	+ save_trace_data_into_json({"traceEvents":[123]}, 'output')
486	mock_dumps.assert_called_once()	486	mock_dumps.assert_called_once()
487	- assert (fake_file.read() == '{"traceEvents":[]}')	487	+ assert (fake_file.read() == '[]')
488	mock_logger.error.assert_called_once()	488	mock_logger.error.assert_called_once()
489		489
490		490