已合并
对齐Pytorch Profiler部分接口 #34217
xfeng创建于 4月23日
对齐Pytorch Profiler部分接口 #34217
已合并
共 6 个文件变更+433-149
| @@ -1,23 +1,27 @@ | |||
| 1 | +# Owner(s): ["oncall: profiler"] | ||
| 1 | import os | 2 | import os |
| 2 | import shutil | 3 | import shutil |
| 3 | from unittest import mock | 4 | from unittest import mock |
| 4 | 5 | ||
| 5 | -import torch | ||
| 6 | - | ||
| 7 | -from torch_npu.npu import Event | ||
| 8 | -from torch_npu.profiler.analysis.prof_common_func._constant import Constant | ||
| 9 | -from torch_npu.profiler.profiler_interface import _ProfInterface | ||
| 10 | -from torch_npu.profiler.profiler_interface import _enable_event_record | ||
| 11 | -from torch_npu.profiler.profiler_interface import _disable_event_record | ||
| 12 | -from torch_npu.profiler._profiler_path_creator import ProfPathCreator | ||
| 13 | -from torch_npu.profiler import supported_activities | ||
| 14 | from torch_npu._C._profiler import ProfilerActivity | 6 | from torch_npu._C._profiler import ProfilerActivity |
| 15 | -from torch_npu.profiler.analysis.prof_common_func._cann_package_manager import CannPackageManager | 7 | +from torch_npu.npu import Event |
| 16 | -from torch_npu.testing.testcase import TestCase, run_tests | 8 | +from torch_npu.profiler import supported_activities |
| 9 | +from torch_npu.profiler._profiler_path_creator import ProfPathCreator | ||
| 10 | +from torch_npu.profiler.analysis.prof_common_func._cann_package_manager import ( | ||
| 11 | + CannPackageManager, | ||
| 12 | +) | ||
| 13 | +from torch_npu.profiler.analysis.prof_common_func._constant import Constant | ||
| 14 | +from torch_npu.profiler.profiler_interface import ( | ||
| 15 | + _disable_event_record, | ||
| 16 | + _enable_event_record, | ||
| 17 | + _ProfInterface, | ||
| 18 | +) | ||
| 19 | +from torch_npu.testing.testcase import run_tests, TestCase | ||
| 20 | + | ||
| 21 | +import torch | ||
| 17 | 22 | ||
| 18 | 23 | ||
| 19 | class TestActionController(TestCase): | 24 | class TestActionController(TestCase): |
| 20 | - | ||
| 21 | 25 | ||
| 22 | def setUpClass(cls): | 26 | def setUpClass(cls): |
| 23 | super().setUpClass() | 27 | super().setUpClass() |
| @@ -34,18 +38,24 @@ class TestActionController(TestCase): | |||
| 34 | ProfPathCreator().init(dir_name=self.prof_dir) | 38 | ProfPathCreator().init(dir_name=self.prof_dir) |
| 35 | 39 | ||
| 36 | def test_init_trace(self): | 40 | def test_init_trace(self): |
| 41 | + self.prof_if.custom_trace_id_callback = lambda: "trace_0" | ||
| 37 | with mock.patch(self.namespace + "._init_profiler") as mock_func: | 42 | with mock.patch(self.namespace + "._init_profiler") as mock_func: |
| 38 | self.prof_if.init_trace() | 43 | self.prof_if.init_trace() |
| 39 | self.assertEqual(1, mock_func.call_count) | 44 | self.assertEqual(1, mock_func.call_count) |
 | |||
| 40 | self.assertTrue(os.path.exists(self.prof_dir)) | 45 | self.assertTrue(os.path.exists(self.prof_dir)) |
| 46 | + self.assertEqual("trace_0", self.prof_if.trace_id) | ||
| 41 | 47 | ||
| 42 | def test_start_trace(self): | 48 | def test_start_trace(self): |
| 43 | - with mock.patch(self.namespace + ".NpuProfilerConfig", return_value="config_obj"), \ | 49 | + with ( |
| 44 | - mock.patch(self.namespace + "._get_syscnt_enable", return_value=True), \ | 50 | + mock.patch( |
| 45 | - mock.patch(self.namespace + "._get_freq", return_value=100), \ | 51 | + self.namespace + ".NpuProfilerConfig", return_value="config_obj" |
| 46 | - mock.patch(self.namespace + "._get_syscnt", return_value=10000), \ | 52 | + ), |
| 47 | - mock.patch(self.namespace + "._get_monotonic", return_value=20000), \ | 53 | + mock.patch(self.namespace + "._get_syscnt_enable", return_value=True), |
| 48 | - mock.patch(self.namespace + "._start_profiler") as mock_start: | 54 | + mock.patch(self.namespace + "._get_freq", return_value=100), |
| 55 | + mock.patch(self.namespace + "._get_syscnt", return_value=10000), | ||
| 56 | + mock.patch(self.namespace + "._get_monotonic", return_value=20000), | ||
| 57 | + mock.patch(self.namespace + "._start_profiler") as mock_start, | ||
| 58 | + ): | ||
| 49 | self.prof_if.start_trace() | 59 | self.prof_if.start_trace() |
| 50 | self.assertEqual(True, self.prof_if.syscnt_enable) | 60 | self.assertEqual(True, self.prof_if.syscnt_enable) |
| 51 | self.assertEqual(100, self.prof_if.freq) | 61 | self.assertEqual(100, self.prof_if.freq) |
| @@ -57,11 +67,13 @@ class TestActionController(TestCase): | |||
| 57 | with mock.patch(self.namespace + "._stop_profiler") as mock_stop: | 67 | with mock.patch(self.namespace + "._stop_profiler") as mock_stop: |
| 58 | self.prof_if.stop_trace() | 68 | self.prof_if.stop_trace() |
| 59 | mock_stop.assert_called_once() | 69 | mock_stop.assert_called_once() |
| 60 | - | 70 | + |
| 61 | def test_finalize_trace(self): | 71 | def test_finalize_trace(self): |
| 62 | - with mock.patch(self.namespace + "._init_profiler"), \ | 72 | + with ( |
| 63 | - mock.patch(self.namespace + "._finalize_profiler") as mock_finalize: | 73 | + mock.patch(self.namespace + "._init_profiler"), |
| 64 | - self.prof_if.metadata = {"key":"val"} | 74 | + mock.patch(self.namespace + "._finalize_profiler") as mock_finalize, |
| 75 | + ): | ||
| 76 | + self.prof_if.metadata = {"key": "val"} | ||
| 65 | self.prof_if.init_trace() | 77 | self.prof_if.init_trace() |
| 66 | self.prof_if.finalize_trace() | 78 | self.prof_if.finalize_trace() |
| 67 | mock_finalize.assert_called_once() | 79 | mock_finalize.assert_called_once() |
| @@ -69,7 +81,9 @@ class TestActionController(TestCase): | |||
| 69 | self.assertTrue(self._check_metadata_json(self.prof_if.prof_path)) | 81 | self.assertTrue(self._check_metadata_json(self.prof_if.prof_path)) |
| 70 | 82 | ||
| 71 | def test_analyse(self): | 83 | def test_analyse(self): |
| 72 | - with mock.patch("torch_npu.profiler.analysis._npu_profiler.NpuProfiler.analyse") as mock_analyse: | 84 | + with mock.patch( |
| 85 | + "torch_npu.profiler.analysis._npu_profiler.NpuProfiler.analyse" | ||
| 86 | + ) as mock_analyse: | ||
| 73 | self.prof_if.analyse() | 87 | self.prof_if.analyse() |
| 74 | mock_analyse.assert_called_once() | 88 | mock_analyse.assert_called_once() |
| 75 | 89 | ||
| @@ -79,6 +93,49 @@ class TestActionController(TestCase): | |||
| 79 | self.assertTrue(ProfilerActivity.CPU in activities) | 93 | self.assertTrue(ProfilerActivity.CPU in activities) |
| 80 | self.assertTrue(ProfilerActivity.NPU in activities) | 94 | self.assertTrue(ProfilerActivity.NPU in activities) |
| 81 | 95 | ||
| 96 | + def test_create_trace_id_should_return_default_when_callback_is_none(self): | ||
| 97 | + self.prof_if.custom_trace_id_callback = None | ||
| 98 | + trace_id = self.prof_if.create_trace_id() | ||
| 99 | + self.assertIsInstance(trace_id, str) | ||
| 100 | + self.assertEqual(32, len(trace_id)) | ||
| 101 | + | ||
| 102 | + def test_create_trace_id_should_return_callback_result_when_callback_is_valid(self): | ||
| 103 | + self.prof_if.custom_trace_id_callback = lambda: "custom_trace_id" | ||
| 104 | + trace_id = self.prof_if.create_trace_id() | ||
| 105 | + self.assertEqual("custom_trace_id", trace_id) | ||
| 106 | + | ||
| 107 | + def test_create_trace_id_should_return_default_when_callback_is_not_callable(self): | ||
| 108 | + self.prof_if.custom_trace_id_callback = "not_callable" | ||
| 109 | + trace_id = self.prof_if.create_trace_id() | ||
| 110 | + self.assertIsInstance(trace_id, str) | ||
| 111 | + self.assertEqual(32, len(trace_id)) | ||
| 112 | + | ||
| 113 | + def test_create_trace_id_should_return_default_when_callback_returns_non_string( | ||
| 114 | + self, | ||
| 115 | + ): | ||
| 116 | + self.prof_if.custom_trace_id_callback = lambda: 12345 | ||
| 117 | + trace_id = self.prof_if.create_trace_id() | ||
| 118 | + self.assertIsInstance(trace_id, str) | ||
| 119 | + self.assertEqual(32, len(trace_id)) | ||
| 120 | + | ||
| 121 | + def test_create_trace_id_should_return_default_when_callback_raises_exception(self): | ||
| 122 | + def bad_callback(): | ||
| 123 | + raise RuntimeError("callback error") | ||
| 124 | + | ||
| 125 | + self.prof_if.custom_trace_id_callback = bad_callback | ||
| 126 | + trace_id = self.prof_if.create_trace_id() | ||
| 127 | + self.assertIsInstance(trace_id, str) | ||
| 128 | + self.assertEqual(32, len(trace_id)) | ||
| 129 | + | ||
| 130 | + def test_create_trace_id_should_return_default_when_callback_returns_too_long_string( | ||
| 131 | + self, | ||
| 132 | + ): | ||
| 133 | + long_str = "a" * (self.prof_if.MAX_TRACE_ID_LEN + 1) | ||
| 134 | + self.prof_if.custom_trace_id_callback = lambda: long_str | ||
| 135 | + trace_id = self.prof_if.create_trace_id() | ||
| 136 | + self.assertIsInstance(trace_id, str) | ||
| 137 | + self.assertEqual(32, len(trace_id)) | ||
| 138 | + | ||
| 82 | def test_event_record_should_have_return_true_attr_when_enable_record(self): | 139 | def test_event_record_should_have_return_true_attr_when_enable_record(self): |
| 83 | _enable_event_record() | 140 | _enable_event_record() |
| 84 | self.assertTrue(hasattr(Event.record, "origin_func")) | 141 | self.assertTrue(hasattr(Event.record, "origin_func")) |
| @@ -96,13 +153,15 @@ class TestActionController(TestCase): | |||
| 96 | def _check_profiler_info_json(self, prof_path: str) -> bool: | 153 | def _check_profiler_info_json(self, prof_path: str) -> bool: |
| 97 | if torch.distributed.is_available() and torch.distributed.is_initialized(): | 154 | if torch.distributed.is_available() and torch.distributed.is_initialized(): |
| 98 | rank_id = torch.distributed.get_rank() | 155 | rank_id = torch.distributed.get_rank() |
| 99 | - path = os.path.join(os.path.realpath(prof_path), f'profiler_info_{rank_id}.json') | 156 | + path = os.path.join( |
| 157 | + os.path.realpath(prof_path), f"profiler_info_{rank_id}.json" | ||
| 158 | + ) | ||
| 100 | else: | 159 | else: |
| 101 | - path = os.path.join(os.path.realpath(prof_path), 'profiler_info.json') | 160 | + path = os.path.join(os.path.realpath(prof_path), "profiler_info.json") |
| 102 | return os.path.exists(path) | 161 | return os.path.exists(path) |
| 103 | 162 | ||
| 104 | def _check_metadata_json(self, prof_path: str) -> bool: | 163 | def _check_metadata_json(self, prof_path: str) -> bool: |
| 105 | - path = os.path.join(os.path.realpath(prof_path), 'profiler_metadata.json') | 164 | + path = os.path.join(os.path.realpath(prof_path), "profiler_metadata.json") |
| 106 | return os.path.exists(path) | 165 | return os.path.exists(path) |
| 107 | 166 | ||
| 108 | def _check_params(self): | 167 | def _check_params(self): |
| @@ -1,48 +1,66 @@ | |||
| 1 | -import torch_npu | 1 | +# Owner(s): ["oncall: profiler"] |
| 2 | -from torch_npu.profiler import ProfilerAction | 2 | +from torch_npu.profiler import ProfilerAction, schedule |
| 3 | -from torch_npu.profiler import schedule | 3 | +from torch_npu.testing.testcase import run_tests, TestCase |
| 4 | -from torch_npu.testing.testcase import TestCase, run_tests | ||
| 5 | 4 | ||
| 6 | 5 | ||
| 7 | class TestScheduler(TestCase): | 6 | class TestScheduler(TestCase): |
| 8 | - | ||
| 9 | def setUp(self): | 7 | def setUp(self): |
| 10 | self.schedule_list = [ | 8 | self.schedule_list = [ |
| 11 | # wait, active, warmup, repeat, skip_first, [step, expected result] | 9 | # wait, active, warmup, repeat, skip_first, [step, expected result] |
| 12 | [ | 10 | [ |
| 13 | - 0, 0, 0, 0, 0, | 11 | + 0, |
| 12 | + 0, | ||
| 13 | + 0, | ||
| 14 | + 0, | ||
| 15 | + 0, | ||
| 14 | [0, ProfilerAction.RECORD_AND_SAVE], | 16 | [0, ProfilerAction.RECORD_AND_SAVE], |
| 15 | [1, ProfilerAction.RECORD_AND_SAVE], | 17 | [1, ProfilerAction.RECORD_AND_SAVE], |
| 16 | - [10, ProfilerAction.RECORD_AND_SAVE] | 18 | + [10, ProfilerAction.RECORD_AND_SAVE], |
| 17 | ], | 19 | ], |
| 18 | [ | 20 | [ |
| 19 | - -1, -1, -1, -1, -1, | 21 | + -1, |
| 22 | + -1, | ||
| 23 | + -1, | ||
| 24 | + -1, | ||
| 25 | + -1, | ||
| 20 | [0, ProfilerAction.RECORD_AND_SAVE], | 26 | [0, ProfilerAction.RECORD_AND_SAVE], |
| 21 | [1, ProfilerAction.RECORD_AND_SAVE], | 27 | [1, ProfilerAction.RECORD_AND_SAVE], |
| 22 | - [10, ProfilerAction.RECORD_AND_SAVE] | 28 | + [10, ProfilerAction.RECORD_AND_SAVE], |
| 23 | ], | 29 | ], |
| 24 | [ | 30 | [ |
| 25 | - 2, 2, 2, 0, 2, | 31 | + 2, |
| 32 | + 2, | ||
| 33 | + 2, | ||
| 34 | + 0, | ||
| 35 | + 2, | ||
| 26 | [0, ProfilerAction.NONE], | 36 | [0, ProfilerAction.NONE], |
| 27 | [1, ProfilerAction.NONE], | 37 | [1, ProfilerAction.NONE], |
| 28 | [3, ProfilerAction.NONE], | 38 | [3, ProfilerAction.NONE], |
| 29 | [5, ProfilerAction.WARMUP], | 39 | [5, ProfilerAction.WARMUP], |
| 30 | [6, ProfilerAction.RECORD], | 40 | [6, ProfilerAction.RECORD], |
| 31 | [7, ProfilerAction.RECORD_AND_SAVE], | 41 | [7, ProfilerAction.RECORD_AND_SAVE], |
| 32 | - [13, ProfilerAction.RECORD_AND_SAVE] | 42 | + [13, ProfilerAction.RECORD_AND_SAVE], |
| 33 | ], | 43 | ], |
| 34 | [ | 44 | [ |
| 35 | - 2, 2, 2, 1, 2, | 45 | + 2, |
| 46 | + 2, | ||
| 47 | + 2, | ||
| 48 | + 1, | ||
| 49 | + 2, | ||
| 36 | [0, ProfilerAction.NONE], | 50 | [0, ProfilerAction.NONE], |
| 37 | [1, ProfilerAction.NONE], | 51 | [1, ProfilerAction.NONE], |
| 38 | [3, ProfilerAction.NONE], | 52 | [3, ProfilerAction.NONE], |
| 39 | [5, ProfilerAction.WARMUP], | 53 | [5, ProfilerAction.WARMUP], |
| 40 | [6, ProfilerAction.RECORD], | 54 | [6, ProfilerAction.RECORD], |
| 41 | [7, ProfilerAction.RECORD_AND_SAVE], | 55 | [7, ProfilerAction.RECORD_AND_SAVE], |
| 42 | - [13, ProfilerAction.NONE] | 56 | + [13, ProfilerAction.NONE], |
| 43 | ], | 57 | ], |
| 44 | [ | 58 | [ |
| 45 | - 2, 2, 2, 4, 2, | 59 | + 2, |
| 60 | + 2, | ||
| 61 | + 2, | ||
| 62 | + 4, | ||
| 63 | + 2, | ||
| 46 | [0, ProfilerAction.NONE], | 64 | [0, ProfilerAction.NONE], |
| 47 | [1, ProfilerAction.NONE], | 65 | [1, ProfilerAction.NONE], |
| 48 | [3, ProfilerAction.NONE], | 66 | [3, ProfilerAction.NONE], |
| @@ -63,14 +81,87 @@ class TestScheduler(TestCase): | |||
| 63 | active=samples[1], | 81 | active=samples[1], |
| 64 | warmup=samples[2], | 82 | warmup=samples[2], |
| 65 | repeat=samples[3], | 83 | repeat=samples[3], |
| 66 | - skip_first=samples[4] | 84 | + skip_first=samples[4], |
| 67 | ) | 85 | ) |
| 68 | - test_pair = samples[self.test_pair_idx:] | 86 | + test_pair = samples[self.test_pair_idx :] |
| 69 | self.assertTrue(sche_inst.active >= 1) | 87 | self.assertTrue(sche_inst.active >= 1) |
| 70 | for step, expect_result in test_pair: | 88 | for step, expect_result in test_pair: |
| 71 | result = sche_inst(step) | 89 | result = sche_inst(step) |
| 72 | self.assertEqual(result, expect_result) | 90 | self.assertEqual(result, expect_result) |
| 73 | 91 | ||
| 92 | + def test_skip_first_wait_should_works_when_non_zero(self): | ||
| 93 | + test_schedule = schedule( | ||
| 94 | + skip_first=1, wait=2, warmup=1, active=2, repeat=2, skip_first_wait=1 | ||
| 95 | + ) | ||
| 96 | + test_schedule_expected_outputs = [ | ||
| 97 | + # repeat No. 1 begin | ||
| 98 | + # skip first 1 | ||
| 99 | + ProfilerAction.NONE, | ||
| 100 | + # warmup 1 | ||
| 101 | + ProfilerAction.WARMUP, | ||
| 102 | + # active 1 begin | ||
| 103 | + ProfilerAction.RECORD, | ||
| 104 | + ProfilerAction.RECORD_AND_SAVE, | ||
| 105 | + # active 1 end | ||
| 106 | + # repeat No. 1 end | ||
| 107 | + # --- | ||
| 108 | + # repeat No. 2 begin | ||
| 109 | + # wait 2 | ||
| 110 | + ProfilerAction.NONE, | ||
| 111 | + ProfilerAction.NONE, | ||
| 112 | + # warmup 1 | ||
| 113 | + ProfilerAction.WARMUP, | ||
| 114 | + # active 2 begin | ||
| 115 | + ProfilerAction.RECORD, | ||
| 116 | + ProfilerAction.RECORD_AND_SAVE, | ||
| 117 | + # active 2 end | ||
| 118 | + # repeat No. 2 end | ||
| 119 | + ProfilerAction.NONE, | ||
| 120 | + ProfilerAction.NONE, | ||
| 121 | + ProfilerAction.NONE, | ||
| 122 | + ProfilerAction.NONE, | ||
| 123 | + ] | ||
| 124 | + for step in range(len(test_schedule_expected_outputs)): | ||
| 125 | + self.assertEqual(test_schedule(step), test_schedule_expected_outputs[step]) | ||
| 126 | + | ||
| 127 | + def test_skip_first_wait_should_be_reset_when_invalid(self): | ||
| 128 | + test_schedule = schedule( | ||
| 129 | + skip_first=1, wait=2, warmup=1, active=2, repeat=2, skip_first_wait=0.5 | ||
| 130 | + ) | ||
| 131 | + test_schedule_expected_outputs = [ | ||
| 132 | + # skip first 1 | ||
| 133 | + ProfilerAction.NONE, | ||
| 134 | + # repeat No. 1 begin | ||
| 135 | + # wait 2 | ||
| 136 | + ProfilerAction.NONE, | ||
| 137 | + ProfilerAction.NONE, | ||
| 138 | + # warmup 1 | ||
| 139 | + ProfilerAction.WARMUP, | ||
| 140 | + # active 1 begin | ||
| 141 | + ProfilerAction.RECORD, | ||
| 142 | + ProfilerAction.RECORD_AND_SAVE, | ||
| 143 | + # active 1 end | ||
| 144 | + # repeat No. 1 end | ||
| 145 | + # --- | ||
| 146 | + # repeat No. 2 begin | ||
| 147 | + # wait 2 | ||
| 148 | + ProfilerAction.NONE, | ||
| 149 | + ProfilerAction.NONE, | ||
| 150 | + # warmup 1 | ||
| 151 | + ProfilerAction.WARMUP, | ||
| 152 | + # active 2 begin | ||
| 153 | + ProfilerAction.RECORD, | ||
| 154 | + ProfilerAction.RECORD_AND_SAVE, | ||
| 155 | + # active 2 end | ||
| 156 | + # repeat No. 2 end | ||
| 157 | + ProfilerAction.NONE, | ||
| 158 | + ProfilerAction.NONE, | ||
| 159 | + ProfilerAction.NONE, | ||
| 160 | + ProfilerAction.NONE, | ||
| 161 | + ] | ||
| 162 | + for step in range(len(test_schedule_expected_outputs)): | ||
| 163 | + self.assertEqual(test_schedule(step), test_schedule_expected_outputs[step]) | ||
| 164 | + | ||
| 74 | 165 | ||
| 75 | if __name__ == "__main__": | 166 | if __name__ == "__main__": |
| 76 | - run_tests() | 167 | + run_tests() |
| @@ -1683,7 +1683,7 @@ | |||
| 1683 | "signature": "(profiler_level: int = 'Level0', aic_metrics: int = 'ACL_AICORE_PIPE_UTILIZATION', l2_cache: bool = False, msprof_tx: bool = False, mstx: bool = False, data_simplification: bool = True, record_op_args: bool = False, op_attr: bool = False, gc_detect_threshold: float = None, export_type: Union[str, list] = 'text', host_sys: list = None, sys_io: bool = False, sys_interconnection: bool = False, mstx_domain_include: list = None, mstx_domain_exclude: list = None)" | 1683 | "signature": "(profiler_level: int = 'Level0', aic_metrics: int = 'ACL_AICORE_PIPE_UTILIZATION', l2_cache: bool = False, msprof_tx: bool = False, mstx: bool = False, data_simplification: bool = True, record_op_args: bool = False, op_attr: bool = False, gc_detect_threshold: float = None, export_type: Union[str, list] = 'text', host_sys: list = None, sys_io: bool = False, sys_interconnection: bool = False, mstx_domain_include: list = None, mstx_domain_exclude: list = None)" |
| 1684 | }, | 1684 | }, |
| 1685 | "torch_npu.profiler.profile": { | 1685 | "torch_npu.profiler.profile": { |
| 1686 | - "signature": "(*, activities: Optional[Iterable[torch_npu._C._profiler.ProfilerActivity]] = None, schedule: Optional[Callable[[int], torch_npu.profiler.scheduler.ProfilerAction]] = None, on_trace_ready: Optional[Callable[..., Any]] = None, record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False, experimental_config: Optional[torch_npu.profiler.experimental_config._ExperimentalConfig] = None, use_cuda: Optional[bool] = None)" | 1686 | + "signature": "(*, activities: Optional[collections.abc.Iterable[torch_npu._C._profiler.ProfilerActivity]] = None, schedule: Optional[collections.abc.Callable[[int], torch_npu.profiler.scheduler.ProfilerAction]] = None, on_trace_ready: Optional[collections.abc.Callable[..., Any]] = None, record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False, experimental_config: Optional[torch_npu.profiler.experimental_config._ExperimentalConfig] = None, custom_trace_id_callback: Optional[collections.abc.Callable[[], str]] = None, use_cuda: Optional[bool] = None)" |
| 1687 | }, | 1687 | }, |
| 1688 | "torch_npu.profiler.profile.start": { | 1688 | "torch_npu.profiler.profile.start": { |
| 1689 | "signature": "(self)" | 1689 | "signature": "(self)" |
| @@ -1694,6 +1694,12 @@ | |||
| 1694 | "torch_npu.profiler.profile.step": { | 1694 | "torch_npu.profiler.profile.step": { |
| 1695 | "signature": "(self)" | 1695 | "signature": "(self)" |
| 1696 | }, | 1696 | }, |
| 1697 | + "torch_npu.profiler.profile.get_trace_id": { | ||
| 1698 | + "signature": "(self) -> str" | ||
| 1699 | + }, | ||
| 1700 | + "torch_npu.profiler.profile.set_custom_trace_id_callback": { | ||
| 1701 | + "signature": "(self, callback: collections.abc.Callable[[], str]) -> None" | ||
| 1702 | + }, | ||
| 1697 | "torch_npu.profiler.profile.enable_profiler_in_child_thread": { | 1703 | "torch_npu.profiler.profile.enable_profiler_in_child_thread": { |
| 1698 | "signature": "(record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False)" | 1704 | "signature": "(record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False)" |
| 1699 | }, | 1705 | }, |
| @@ -1701,7 +1707,7 @@ | |||
| 1701 | "signature": "()" | 1707 | "signature": "()" |
| 1702 | }, | 1708 | }, |
| 1703 | "torch_npu.profiler.schedule": { | 1709 | "torch_npu.profiler.schedule": { |
| 1704 | - "signature": "(wait: int, active: int, warmup: int = 0, repeat: int = 0, skip_first: int = 0) -> None" | 1710 | + "signature": "(wait: int, active: int, warmup: int = 0, repeat: int = 0, skip_first: int = 0, skip_first_wait: int = 0) -> None" |
| 1705 | }, | 1711 | }, |
| 1706 | "torch_npu.profiler.supported_activities": { | 1712 | "torch_npu.profiler.supported_activities": { |
| 1707 | "signature": "()" | 1713 | "signature": "()" |
| @@ -1716,7 +1722,7 @@ | |||
| 1716 | "signature": "()" | 1722 | "signature": "()" |
| 1717 | }, | 1723 | }, |
| 1718 | "torch_npu.profiler.tensorboard_trace_handler": { | 1724 | "torch_npu.profiler.tensorboard_trace_handler": { |
| 1719 | - "signature": "(dir_name: str = None, worker_name: str = None, analyse_flag: bool = True, async_mode: bool = False)" | 1725 | + "signature": "(dir_name: Optional[str] = None, worker_name: Optional[str] = None, analyse_flag: bool = True, async_mode: bool = False)" |
| 1720 | }, | 1726 | }, |
| 1721 | "torch_npu.profiler.dynamic_profile.init": { | 1727 | "torch_npu.profiler.dynamic_profile.init": { |
| 1722 | "signature": "(path: str)" | 1728 | "signature": "(path: str)" |
| @@ -1749,10 +1755,10 @@ | |||
| 1749 | "signature": "()" | 1755 | "signature": "()" |
| 1750 | }, | 1756 | }, |
| 1751 | "torch_npu.profiler.profiler.analyse": { | 1757 | "torch_npu.profiler.profiler.analyse": { |
| 1752 | - "signature": "(profiler_path: str, max_process_number: int = 96, export_type: Union[str, list] = None)" | 1758 | + "signature": "(profiler_path: str, max_process_number: int = 320, export_type: Union[str, list, NoneType] = None)" |
| 1753 | }, | 1759 | }, |
| 1754 | "torch_npu.profiler.profiler.profile": { | 1760 | "torch_npu.profiler.profiler.profile": { |
| 1755 | - "signature": "(*, activities: Optional[Iterable[torch_npu._C._profiler.ProfilerActivity]] = None, schedule: Optional[Callable[[int], torch_npu.profiler.scheduler.ProfilerAction]] = None, on_trace_ready: Optional[Callable[..., Any]] = None, record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False, experimental_config: Optional[torch_npu.profiler.experimental_config._ExperimentalConfig] = None, use_cuda: Optional[bool] = None)" | 1761 | + "signature": "(*, activities: Optional[collections.abc.Iterable[torch_npu._C._profiler.ProfilerActivity]] = None, schedule: Optional[collections.abc.Callable[[int], torch_npu.profiler.scheduler.ProfilerAction]] = None, on_trace_ready: Optional[collections.abc.Callable[..., Any]] = None, record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False, experimental_config: Optional[torch_npu.profiler.experimental_config._ExperimentalConfig] = None, custom_trace_id_callback: Optional[collections.abc.Callable[[], str]] = None, use_cuda: Optional[bool] = None)" |
| 1756 | }, | 1762 | }, |
| 1757 | "torch_npu.profiler.profiler.profile.start": { | 1763 | "torch_npu.profiler.profiler.profile.start": { |
| 1758 | "signature": "(self)" | 1764 | "signature": "(self)" |
| @@ -1763,6 +1769,12 @@ | |||
| 1763 | "torch_npu.profiler.profiler.profile.step": { | 1769 | "torch_npu.profiler.profiler.profile.step": { |
| 1764 | "signature": "(self)" | 1770 | "signature": "(self)" |
| 1765 | }, | 1771 | }, |
| 1772 | + "torch_npu.profiler.profiler.profile.get_trace_id": { | ||
| 1773 | + "signature": "(self) -> str" | ||
| 1774 | + }, | ||
| 1775 | + "torch_npu.profiler.profiler.profile.set_custom_trace_id_callback": { | ||
| 1776 | + "signature": "(self, callback: collections.abc.Callable[[], str]) -> None" | ||
| 1777 | + }, | ||
| 1766 | "torch_npu.profiler.profiler.profile.enable_profiler_in_child_thread": { | 1778 | "torch_npu.profiler.profiler.profile.enable_profiler_in_child_thread": { |
| 1767 | "signature": "(record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False)" | 1779 | "signature": "(record_shapes: bool = False, profile_memory: bool = False, with_stack: bool = False, with_flops: bool = False, with_modules: bool = False)" |
| 1768 | }, | 1780 | }, |
| @@ -1773,7 +1785,7 @@ | |||
| 1773 | "signature": "()" | 1785 | "signature": "()" |
| 1774 | }, | 1786 | }, |
| 1775 | "torch_npu.profiler.profiler.tensorboard_trace_handler": { | 1787 | "torch_npu.profiler.profiler.tensorboard_trace_handler": { |
| 1776 | - "signature": "(dir_name: str = None, worker_name: str = None, analyse_flag: bool = True, async_mode: bool = False)" | 1788 | + "signature": "(dir_name: Optional[str] = None, worker_name: Optional[str] = None, analyse_flag: bool = True, async_mode: bool = False)" |
| 1777 | }, | 1789 | }, |
| 1778 | "torch_npu.profiler.profiler_interface.supported_activities": { | 1790 | "torch_npu.profiler.profiler_interface.supported_activities": { |
| 1779 | "signature": "()" | 1791 | "signature": "()" |
| @@ -1782,7 +1794,7 @@ | |||
| 1782 | "signature": "(value, names=None, *, module=None, qualname=None, type=None, start=1)" | 1794 | "signature": "(value, names=None, *, module=None, qualname=None, type=None, start=1)" |
| 1783 | }, | 1795 | }, |
| 1784 | "torch_npu.profiler.scheduler.Schedule": { | 1796 | "torch_npu.profiler.scheduler.Schedule": { |
| 1785 | - "signature": "(wait: int, active: int, warmup: int = 0, repeat: int = 0, skip_first: int = 0) -> None" | 1797 | + "signature": "(wait: int, active: int, warmup: int = 0, repeat: int = 0, skip_first: int = 0, skip_first_wait: int = 0) -> None" |
| 1786 | }, | 1798 | }, |
| 1787 | "torch_npu.testing.common_distributed.TestSkip": { | 1799 | "torch_npu.testing.common_distributed.TestSkip": { |
| 1788 | "signature": "(exit_code, message)" | 1800 | "signature": "(exit_code, message)" |
| @@ -1,50 +1,48 @@ | |||
| 1 | -import os.path | 1 | +# ruff: noqa: UP045, UP007 |
| 2 | import json | 2 | import json |
| 3 | +import os.path | ||
| 4 | +from collections.abc import Callable, Iterable | ||
| 3 | from sys import getsizeof | 5 | from sys import getsizeof |
| 4 | -from typing import Optional, Iterable, Callable, Any, Union | 6 | +from typing import Any, Optional, Union |
| 5 | 7 | ||
| 6 | import torch.autograd.profiler as prof | 8 | import torch.autograd.profiler as prof |
| 7 | import torch_npu.npu | 9 | import torch_npu.npu |
| 8 | from torch_npu._C._profiler import ( | 10 | from torch_npu._C._profiler import ( |
| 9 | - _enable_profiler_in_child_thread, | ||
| 10 | _disable_profiler_in_child_thread, | 11 | _disable_profiler_in_child_thread, |
| 12 | + _enable_profiler_in_child_thread, | ||
| 11 | _ExperimentalConfig as C_ExperimentalConfig, | 13 | _ExperimentalConfig as C_ExperimentalConfig, |
| 14 | + NpuProfilerConfig, | ||
| 12 | ProfilerActivity, | 15 | ProfilerActivity, |
| 13 | - NpuProfilerConfig | ||
| 14 | ) | 16 | ) |
| 15 | from torch_npu.utils._error_code import ErrCode, prof_error | 17 | from torch_npu.utils._error_code import ErrCode, prof_error |
| 16 | 18 | ||
| 17 | -from .experimental_config import _ExperimentalConfig | ||
| 18 | -from ._profiler_path_creator import ProfPathCreator | ||
| 19 | -from .profiler_interface import _ProfInterface, supported_activities | ||
| 20 | -from ._profiler_action_controller import ProfActionController | ||
| 21 | -from .scheduler import _default_schedule_fn, ProfilerAction | ||
| 22 | -from .analysis.prof_common_func._constant import Constant | ||
| 23 | -from .analysis.prof_common_func._constant import print_warn_msg | ||
| 24 | -from .analysis.prof_common_func._utils import no_exception_func | ||
| 25 | -from .analysis._npu_profiler import NpuProfiler | ||
| 26 | -from .analysis.prof_common_func._path_manager import ProfilerPathManager | ||
| 27 | from ..utils._path_manager import PathManager | 19 | from ..utils._path_manager import PathManager |
| 20 | +from ._profiler_action_controller import ProfActionController | ||
| 21 | +from ._profiler_path_creator import ProfPathCreator | ||
| 22 | +from .analysis._npu_profiler import NpuProfiler | ||
| 23 | +from .analysis.prof_common_func._constant import Constant, print_warn_msg | ||
| 24 | +from .analysis.prof_common_func._path_manager import ProfilerPathManager | ||
| 25 | +from .analysis.prof_common_func._utils import no_exception_func | ||
| 26 | +from .experimental_config import _ExperimentalConfig | ||
| 27 | +from .profiler_interface import _ProfInterface, supported_activities | ||
| 28 | +from .scheduler import _default_schedule_fn, ProfilerAction | ||
| 28 | 29 | ||
| 29 | -__all__ = [ | 30 | + |
| 30 | - 'supported_activities', | 31 | +__all__ = ["supported_activities", "analyse", "tensorboard_trace_handler", "profile"] |
| 31 | - 'analyse', | ||
| 32 | - 'tensorboard_trace_handler', | ||
| 33 | - 'profile' | ||
| 34 | -] | ||
| 35 | 32 | ||
| 36 | 33 | ||
| 37 | class _KinetoProfile: | 34 | class _KinetoProfile: |
| 38 | def __init__( | 35 | def __init__( |
| 39 | - self, | 36 | + self, |
| 40 | - *, | 37 | + *, |
| 41 | - activities: Optional[Iterable[ProfilerActivity]] = None, | 38 | + activities: Optional[Iterable[ProfilerActivity]] = None, |
| 42 | - record_shapes: bool = False, | 39 | + record_shapes: bool = False, |
| 43 | - profile_memory: bool = False, | 40 | + profile_memory: bool = False, |
| 44 | - with_stack: bool = False, | 41 | + with_stack: bool = False, |
| 45 | - with_flops: bool = False, | 42 | + with_flops: bool = False, |
| 46 | - with_modules: bool = False, | 43 | + with_modules: bool = False, |
| 47 | - experimental_config: Optional[_ExperimentalConfig] = None, | 44 | + experimental_config: Optional[_ExperimentalConfig] = None, |
| 45 | + custom_trace_id_callback: Optional[Callable[[], str]] = None, | ||
| 48 | ): | 46 | ): |
| 49 | self.metadata = {} | 47 | self.metadata = {} |
| 50 | self.prof_if = _ProfInterface( | 48 | self.prof_if = _ProfInterface( |
| @@ -55,7 +53,8 @@ class _KinetoProfile: | |||
| 55 | with_flops=with_flops, | 53 | with_flops=with_flops, |
| 56 | with_modules=with_modules, | 54 | with_modules=with_modules, |
| 57 | experimental_config=experimental_config, | 55 | experimental_config=experimental_config, |
| 58 | - metadata=self.metadata | 56 | + metadata=self.metadata, |
| 57 | + custom_trace_id_callback=custom_trace_id_callback, | ||
| 59 | ) | 58 | ) |
| 60 | self.max_meta_size = 50 * 1024 | 59 | self.max_meta_size = 50 * 1024 |
| 61 | self.max_str_len = 4096 | 60 | self.max_str_len = 4096 |
| @@ -81,7 +80,10 @@ class _KinetoProfile: | |||
| 81 | PathManager.check_input_file_path(output_path) | 80 | PathManager.check_input_file_path(output_path) |
| 82 | file_name = os.path.basename(output_path) | 81 | file_name = os.path.basename(output_path) |
| 83 | if not file_name.endswith(".json"): | 82 | if not file_name.endswith(".json"): |
| 84 | - raise RuntimeError("Invalid parameter output_path, which must be a json file." + prof_error(ErrCode.VALUE)) | 83 | + raise RuntimeError( |
| 84 | + "Invalid parameter output_path, which must be a json file." | ||
| 85 | + + prof_error(ErrCode.VALUE) | ||
| 86 | + ) | ||
| 85 | if not self.prof_if.prof_path: | 87 | if not self.prof_if.prof_path: |
| 86 | print_warn_msg("Invalid profiling path.") | 88 | print_warn_msg("Invalid profiling path.") |
| 87 | return | 89 | return |
| @@ -90,14 +92,16 @@ class _KinetoProfile: | |||
| 90 | 92 | ||
| 91 | def add_metadata(self, key: str, value: str): | 93 | def add_metadata(self, key: str, value: str): |
| 92 | if not isinstance(key, str) or not isinstance(value, str): | 94 | if not isinstance(key, str) or not isinstance(value, str): |
| 93 | - print_warn_msg("The key and value of metadata must be string. Skip this metadata.") | 95 | + print_warn_msg( |
| 96 | + "The key and value of metadata must be string. Skip this metadata." | ||
| 97 | + ) | ||
| 94 | return | 98 | return |
| 95 | if not self._check_str_valid(key) or not self._check_str_valid(value): | 99 | if not self._check_str_valid(key) or not self._check_str_valid(value): |
| 96 | print_warn_msg("Invalid input key or value. Skip this metadata.") | 100 | print_warn_msg("Invalid input key or value. Skip this metadata.") |
| 97 | return | 101 | return |
| 98 | add_size = getsizeof(key) + getsizeof(value) | 102 | add_size = getsizeof(key) + getsizeof(value) |
| 99 | if getsizeof(self.metadata) + add_size < self.max_meta_size: | 103 | if getsizeof(self.metadata) + add_size < self.max_meta_size: |
| 100 | - if key in self.metadata.keys(): | 104 | + if key in self.metadata: |
| 101 | print_warn_msg(f"{key} is already saved as metadata, override it.") | 105 | print_warn_msg(f"{key} is already saved as metadata, override it.") |
| 102 | self.metadata[key] = value | 106 | self.metadata[key] = value |
| 103 | else: | 107 | else: |
| @@ -106,7 +110,9 @@ class _KinetoProfile: | |||
| 106 | 110 | ||
| 107 | def add_metadata_json(self, key: str, value: str): | 111 | def add_metadata_json(self, key: str, value: str): |
| 108 | if not isinstance(key, str) or not isinstance(value, str): | 112 | if not isinstance(key, str) or not isinstance(value, str): |
| 109 | - print_warn_msg("The key and value of metadata must be string. Skip this metadata.") | 113 | + print_warn_msg( |
| 114 | + "The key and value of metadata must be string. Skip this metadata." | ||
| 115 | + ) | ||
| 110 | return | 116 | return |
| 111 | if not self._check_str_valid(key) or not self._check_str_valid(value): | 117 | if not self._check_str_valid(key) or not self._check_str_valid(value): |
| 112 | print_warn_msg("Invalid input key or value. Skip this metadata.") | 118 | print_warn_msg("Invalid input key or value. Skip this metadata.") |
| @@ -114,11 +120,13 @@ class _KinetoProfile: | |||
| 114 | add_size = getsizeof(key) + getsizeof(value) | 120 | add_size = getsizeof(key) + getsizeof(value) |
| 115 | if getsizeof(self.metadata) + add_size < self.max_meta_size: | 121 | if getsizeof(self.metadata) + add_size < self.max_meta_size: |
| 116 | try: | 122 | try: |
| 117 | - if key in self.metadata.keys(): | 123 | + if key in self.metadata: |
| 118 | print_warn_msg(f"{key} is already saved as metadata, override it.") | 124 | print_warn_msg(f"{key} is already saved as metadata, override it.") |
| 119 | self.metadata[key] = json.loads(value) | 125 | self.metadata[key] = json.loads(value) |
| 120 | except ValueError: | 126 | except ValueError: |
| 121 | - print_warn_msg("The metadata value must be json format string. Skip this metadata") | 127 | + print_warn_msg( |
| 128 | + "The metadata value must be json format string. Skip this metadata" | ||
| 129 | + ) | ||
| 122 | else: | 130 | else: |
| 123 | print_warn_msg("Too many metadata added. Skip this metadata") | 131 | print_warn_msg("Too many metadata added. Skip this metadata") |
| 124 | 132 | ||
| @@ -130,8 +138,10 @@ class _KinetoProfile: | |||
| 130 | print_warn_msg("Function export_stacks() requires with_stack=True.") | 138 | print_warn_msg("Function export_stacks() requires with_stack=True.") |
| 131 | return | 139 | return |
| 132 | if metric not in self._support_export_stacks_metrics(): | 140 | if metric not in self._support_export_stacks_metrics(): |
| 133 | - print_warn_msg("Metric should be self_cpu_time_total or self_npu_time_total." | 141 | + print_warn_msg( |
| 134 | - "Here it is presumed to be self_cpu_time_total.") | 142 | + "Metric should be self_cpu_time_total or self_npu_time_total." |
| 143 | + "Here it is presumed to be self_cpu_time_total." | ||
| 144 | + ) | ||
| 135 | metric = Constant.METRIC_CPU_TIME | 145 | metric = Constant.METRIC_CPU_TIME |
| 136 | if not self.prof_if.prof_path: | 146 | if not self.prof_if.prof_path: |
| 137 | print_warn_msg("Invalid profiling path.") | 147 | print_warn_msg("Invalid profiling path.") |
| @@ -139,7 +149,9 @@ class _KinetoProfile: | |||
| 139 | self.prof_if.analyse(Constant.EXPORT_STACK, output_path, metric=metric) | 149 | self.prof_if.analyse(Constant.EXPORT_STACK, output_path, metric=metric) |
| 140 | 150 | ||
| 141 | 151 | ||
| 142 | - def export_memory_timeline(self, output_path: str, device: Optional[str] = None) -> None: | 152 | + def export_memory_timeline( |
| 153 | + self, output_path: str, device: Optional[str] = None | ||
| 154 | + ) -> None: | ||
| 143 | if device is None: | 155 | if device is None: |
| 144 | device = "npu:0" if torch_npu.npu.is_available() else "cpu" | 156 | device = "npu:0" if torch_npu.npu.is_available() else "cpu" |
| 145 | missing = [] | 157 | missing = [] |
| @@ -156,7 +168,9 @@ class _KinetoProfile: | |||
| 156 | if not self.prof_if.prof_path: | 168 | if not self.prof_if.prof_path: |
| 157 | print_warn_msg("Invalid profiling path.") | 169 | print_warn_msg("Invalid profiling path.") |
| 158 | return | 170 | return |
| 159 | - self.prof_if.analyse(Constant.EXPORT_MEMORY_TIMELINE, output_path, device=device) | 171 | + self.prof_if.analyse( |
| 172 | + Constant.EXPORT_MEMORY_TIMELINE, output_path, device=device | ||
| 173 | + ) | ||
| 160 | 174 | ||
| 161 | def _check_str_valid(self, input_str: str): | 175 | def _check_str_valid(self, input_str: str): |
| 162 | if len(input_str) > self.max_str_len: | 176 | if len(input_str) > self.max_str_len: |
| @@ -168,8 +182,12 @@ class _KinetoProfile: | |||
| 168 | 182 | ||
| 169 | 183 | ||
| 170 | 184 | ||
| 171 | -def tensorboard_trace_handler(dir_name: str = None, worker_name: str = None, | 185 | +def tensorboard_trace_handler( |
| 172 | - analyse_flag: bool = True, async_mode: bool = False): | 186 | + dir_name: Optional[str] = None, |
| 187 | + worker_name: Optional[str] = None, | ||
| 188 | + analyse_flag: bool = True, | ||
| 189 | + async_mode: bool = False, | ||
| 190 | +): | ||
| 173 | ProfPathCreator().init(worker_name=worker_name, dir_name=dir_name) | 191 | ProfPathCreator().init(worker_name=worker_name, dir_name=dir_name) |
| 174 | if not isinstance(analyse_flag, bool): | 192 | if not isinstance(analyse_flag, bool): |
| 175 | print_warn_msg("analyse_flag is not bool, set by default.") | 193 | print_warn_msg("analyse_flag is not bool, set by default.") |
| @@ -187,27 +205,30 @@ def tensorboard_trace_handler(dir_name: str = None, worker_name: str = None, | |||
| 187 | 205 | ||
| 188 | class profile(_KinetoProfile): | 206 | class profile(_KinetoProfile): |
| 189 | def __init__( | 207 | def __init__( |
| 190 | - self, | 208 | + self, |
| 191 | - *, | 209 | + *, |
| 192 | - activities: Optional[Iterable[ProfilerActivity]] = None, | 210 | + activities: Optional[Iterable[ProfilerActivity]] = None, |
| 193 | - schedule: Optional[Callable[[int], ProfilerAction]] = None, | 211 | + schedule: Optional[Callable[[int], ProfilerAction]] = None, |
| 194 | - on_trace_ready: Optional[Callable[..., Any]] = None, | 212 | + on_trace_ready: Optional[Callable[..., Any]] = None, |
| 195 | - record_shapes: bool = False, | 213 | + record_shapes: bool = False, |
| 196 | - profile_memory: bool = False, | 214 | + profile_memory: bool = False, |
| 197 | - with_stack: bool = False, | 215 | + with_stack: bool = False, |
| 198 | - with_flops: bool = False, | 216 | + with_flops: bool = False, |
| 199 | - with_modules: bool = False, | 217 | + with_modules: bool = False, |
| 200 | - experimental_config: Optional[_ExperimentalConfig] = None, | 218 | + experimental_config: Optional[_ExperimentalConfig] = None, |
| 201 | - # deprecated: | 219 | + custom_trace_id_callback: Optional[Callable[[], str]] = None, |
| 202 | - use_cuda: Optional[bool] = None, | 220 | + # deprecated: |
| 221 | + use_cuda: Optional[bool] = None, | ||
| 203 | ): | 222 | ): |
| 204 | - super().__init__(activities=activities, | 223 | + super().__init__( |
| 205 | - record_shapes=record_shapes, | 224 | + activities=activities, |
| 206 | - profile_memory=profile_memory, | 225 | + record_shapes=record_shapes, |
| 207 | - with_stack=with_stack, | 226 | + profile_memory=profile_memory, |
| 208 | - with_flops=with_flops, | 227 | + with_stack=with_stack, |
| 209 | - with_modules=with_modules, | 228 | + with_flops=with_flops, |
| 210 | - experimental_config=experimental_config) | 229 | + with_modules=with_modules, |
| 230 | + experimental_config=experimental_config, | ||
| 231 | + ) | ||
| 211 | activities_set = set(activities) if activities else supported_activities() | 232 | activities_set = set(activities) if activities else supported_activities() |
| 212 | if schedule and isinstance(schedule, Callable): | 233 | if schedule and isinstance(schedule, Callable): |
| 213 | self.schedule = schedule | 234 | self.schedule = schedule |
| @@ -230,7 +251,8 @@ class profile(_KinetoProfile): | |||
| 230 | with_modules=with_modules, | 251 | with_modules=with_modules, |
| 231 | experimental_config=experimental_config, | 252 | experimental_config=experimental_config, |
| 232 | schedule=self.schedule, | 253 | schedule=self.schedule, |
| 233 | - metadata=self.metadata | 254 | + metadata=self.metadata, |
| 255 | + custom_trace_id_callback=custom_trace_id_callback, | ||
| 234 | ) | 256 | ) |
| 235 | self.on_trace_ready = on_trace_ready | 257 | self.on_trace_ready = on_trace_ready |
| 236 | self.step_num = 0 | 258 | self.step_num = 0 |
| @@ -240,7 +262,9 @@ class profile(_KinetoProfile): | |||
| 240 | if use_cuda is not None: | 262 | if use_cuda is not None: |
| 241 | print_warn_msg("This is npu environment, use_cuda is invalid") | 263 | print_warn_msg("This is npu environment, use_cuda is invalid") |
| 242 | self.stopped = False | 264 | self.stopped = False |
| 243 | - self.action_controller = ProfActionController(self, self.prof_if, self.on_trace_ready) | 265 | + self.action_controller = ProfActionController( |
| 266 | + self, self.prof_if, self.on_trace_ready | ||
| 267 | + ) | ||
| 244 | 268 | ||
| 245 | 269 | ||
| 246 | def __enter__(self): | 270 | def __enter__(self): |
| @@ -253,7 +277,7 @@ class profile(_KinetoProfile): | |||
| 253 | 277 | ||
| 254 | 278 | ||
| 255 | def __del__(self): | 279 | def __del__(self): |
| 256 | - if self.stopped == False: | 280 | + if not self.stopped: |
| 257 | self.stop() | 281 | self.stop() |
| 258 | 282 | ||
| 259 | 283 | ||
| @@ -267,7 +291,9 @@ class profile(_KinetoProfile): | |||
| 267 | ProfPathCreator().init(export_only_mode=True) | 291 | ProfPathCreator().init(export_only_mode=True) |
| 268 | self.action_controller.transit_action(ProfilerAction.NONE, self.current_action) | 292 | self.action_controller.transit_action(ProfilerAction.NONE, self.current_action) |
| 269 | if self.record_steps: | 293 | if self.record_steps: |
| 270 | - self.step_rec_fn = prof.record_function("ProfilerStep#" + str(self.step_num + self._step_num_offset)) | 294 | + self.step_rec_fn = prof.record_function( |
| 295 | + "ProfilerStep#" + str(self.step_num + self._step_num_offset) | ||
| 296 | + ) | ||
| 271 | self.step_rec_fn.__enter__() | 297 | self.step_rec_fn.__enter__() |
| 272 | 298 | ||
| 273 | 299 | ||
| @@ -289,31 +315,51 @@ class profile(_KinetoProfile): | |||
| 289 | self.current_action = self.schedule(self.step_num) | 315 | self.current_action = self.schedule(self.step_num) |
| 290 | self.action_controller.transit_action(prev_action, self.current_action) | 316 | self.action_controller.transit_action(prev_action, self.current_action) |
| 291 | if self.record_steps: | 317 | if self.record_steps: |
| 292 | - self.step_rec_fn = prof.record_function("ProfilerStep#" + str(self.step_num + self._step_num_offset)) | 318 | + self.step_rec_fn = prof.record_function( |
| 319 | + "ProfilerStep#" + str(self.step_num + self._step_num_offset) | ||
| 320 | + ) | ||
| 293 | self.step_rec_fn.__enter__() | 321 | self.step_rec_fn.__enter__() |
| 294 | 322 | ||
| 323 | + | ||
| 324 | + def set_custom_trace_id_callback(self, callback: Callable[[], str]) -> None: | ||
🚨 set_custom_trace_id_callback 不会更新已生成的 trace_id
参考依据
   xfeng 4月24日 评论： 4月24日 评论： | |||
| 325 | + self.prof_if.custom_trace_id_callback = callback | ||
| 326 | + | ||
| 327 | + | ||
| 328 | + def get_trace_id(self) -> str: | ||
| 329 | + return self.prof_if.trace_id | ||
| 330 | + | ||
| 295 | 331 | ||
| 296 | 332 | ||
| 297 | - def enable_profiler_in_child_thread(cls, | 333 | + def enable_profiler_in_child_thread( |
| 298 | - record_shapes: bool = False, | 334 | + cls, |
| 299 | - profile_memory: bool = False, | 335 | + record_shapes: bool = False, |
| 300 | - with_stack: bool = False, | 336 | + profile_memory: bool = False, |
| 301 | - with_flops: bool = False, | 337 | + with_stack: bool = False, |
| 302 | - with_modules: bool = False): | 338 | + with_flops: bool = False, |
| 339 | + with_modules: bool = False, | ||
| 340 | + ): | ||
| 303 | params = { | 341 | params = { |
| 304 | - 'record_shapes': record_shapes, | 342 | + "record_shapes": record_shapes, |
| 305 | - 'profile_memory': profile_memory, | 343 | + "profile_memory": profile_memory, |
| 306 | - 'with_stack': with_stack, | 344 | + "with_stack": with_stack, |
| 307 | - 'with_flops': with_flops, | 345 | + "with_flops": with_flops, |
| 308 | - 'with_modules': with_modules | 346 | + "with_modules": with_modules, |
| 309 | } | 347 | } |
| 310 | for param_name, param_value in params.items(): | 348 | for param_name, param_value in params.items(): |
| 311 | if not isinstance(param_value, bool): | 349 | if not isinstance(param_value, bool): |
| 312 | - print_warn_msg(f"{param_name} in enable_profiler_in_child_thread is not bool, reset it to False.") | 350 | + print_warn_msg( |
| 351 | + f"{param_name} in enable_profiler_in_child_thread is not bool, reset it to False." | ||
| 352 | + ) | ||
| 313 | params[param_name] = False | 353 | params[param_name] = False |
| 314 | - npu_prof_config = NpuProfilerConfig('', params['record_shapes'], params['profile_memory'], | 354 | + npu_prof_config = NpuProfilerConfig( |
| 315 | - params['with_stack'], params['with_flops'], params['with_modules'], | 355 | + "", |
| 316 | - C_ExperimentalConfig()) | 356 | + params["record_shapes"], |
| 357 | + params["profile_memory"], | ||
| 358 | + params["with_stack"], | ||
| 359 | + params["with_flops"], | ||
| 360 | + params["with_modules"], | ||
| 361 | + C_ExperimentalConfig(), | ||
| 362 | + ) | ||
| 317 | _enable_profiler_in_child_thread(npu_prof_config) | 363 | _enable_profiler_in_child_thread(npu_prof_config) |
| 318 | 364 | ||
| 319 | 365 | ||
| @@ -324,24 +370,37 @@ class profile(_KinetoProfile): | |||
| 324 | 370 | ||
| 325 | 371 | ||
| 326 | 372 | ||
| 327 | -def analyse(profiler_path: str, max_process_number: int = Constant.DEFAULT_PROCESS_NUMBER, | 373 | +def analyse( |
| 328 | - export_type: Union[str, list] = None): | 374 | + profiler_path: str, |
| 375 | + max_process_number: int = Constant.DEFAULT_PROCESS_NUMBER, | ||
| 376 | + export_type: Optional[Union[str, list]] = None, | ||
| 377 | +): | ||
| 329 | if not isinstance(max_process_number, int) or max_process_number <= 0: | 378 | if not isinstance(max_process_number, int) or max_process_number <= 0: |
| 330 | max_process_number = Constant.DEFAULT_PROCESS_NUMBER | 379 | max_process_number = Constant.DEFAULT_PROCESS_NUMBER |
| 331 | print_warn_msg("Invalid max_process_number, reset it to default!") | 380 | print_warn_msg("Invalid max_process_number, reset it to default!") |
| 332 | if max_process_number > os.cpu_count(): | 381 | if max_process_number > os.cpu_count(): |
| 333 | max_process_number = os.cpu_count() | 382 | max_process_number = os.cpu_count() |
| 334 | - print_warn_msg("max_process_number exceeds the number of cpu cores, reset it to the number of cpu cores!") | 383 | + print_warn_msg( |
| 384 | + "max_process_number exceeds the number of cpu cores, reset it to the number of cpu cores!" | ||
| 385 | + ) | ||
| 335 | if export_type is not None: | 386 | if export_type is not None: |
| 336 | if isinstance(export_type, str): | 387 | if isinstance(export_type, str): |
| 337 | export_type = [export_type] | 388 | export_type = [export_type] |
| 338 | elif isinstance(export_type, list): | 389 | elif isinstance(export_type, list): |
| 339 | export_type = list(set(export_type)) | 390 | export_type = list(set(export_type)) |
| 340 | else: | 391 | else: |
| 341 | - print_warn_msg(f"Invalid parameter export_type: {export_type}, reset it to None.") | 392 | + print_warn_msg( |
| 393 | + f"Invalid parameter export_type: {export_type}, reset it to None." | ||
| 394 | + ) | ||
| 342 | export_type = None | 395 | export_type = None |
| 343 | if export_type is not None: | 396 | if export_type is not None: |
| 344 | - if not export_type or not all(_type in [Constant.Text, Constant.Db] for _type in export_type): | 397 | + if not export_type or not all( |
| 345 | - print_warn_msg(f"Invalid parameter export_type: {export_type}, reset it to None.") | 398 | + _type in [Constant.Text, Constant.Db] for _type in export_type |
| 399 | + ): | ||
| 400 | + print_warn_msg( | ||
| 401 | + f"Invalid parameter export_type: {export_type}, reset it to None." | ||
| 402 | + ) | ||
| 346 | export_type = None | 403 | export_type = None |
| 347 | - NpuProfiler.analyse(profiler_path, max_process_number=max_process_number, export_type=export_type) | 404 | + NpuProfiler.analyse( |
| 405 | + profiler_path, max_process_number=max_process_number, export_type=export_type | ||
| 406 | + ) | ||
| @@ -2,6 +2,7 @@ | |||
| 2 | import functools | 2 | import functools |
| 3 | import os | 3 | import os |
| 4 | import time | 4 | import time |
| 5 | +import uuid | ||
| 5 | from collections.abc import Callable, Iterable | 6 | from collections.abc import Callable, Iterable |
| 6 | from typing import Optional | 7 | from typing import Optional |
| 7 | 8 | ||
| @@ -70,18 +71,21 @@ def _disable_event_record(): | |||
| 70 | 71 | ||
| 71 | class _ProfInterface: | 72 | class _ProfInterface: |
| 72 | PARALLEL_GROUP_KEY = "parallel_group_info" | 73 | PARALLEL_GROUP_KEY = "parallel_group_info" |
| 74 | + TRACE_ID_KEY = "trace_id" | ||
| 75 | + MAX_TRACE_ID_LEN = 1024 | ||
| 73 | 76 | ||
| 74 | def __init__( | 77 | def __init__( |
| 75 | self, | 78 | self, |
| 76 | - activities: Iterable[ProfilerActivity] | None = None, | 79 | + activities: Optional[Iterable[ProfilerActivity]] = None, |
| 77 | record_shapes: bool = False, | 80 | record_shapes: bool = False, |
| 78 | profile_memory: bool = False, | 81 | profile_memory: bool = False, |
| 79 | with_stack: bool = False, | 82 | with_stack: bool = False, |
| 80 | with_flops: bool = False, | 83 | with_flops: bool = False, |
| 81 | with_modules: bool = False, | 84 | with_modules: bool = False, |
| 82 | - schedule: Callable[[int], ProfilerAction] | None = None, | 85 | + schedule: Optional[Callable[[int], ProfilerAction]] = None, |
| 83 | metadata: Optional[dict] = None, | 86 | metadata: Optional[dict] = None, |
| 84 | - experimental_config: _ExperimentalConfig | None = None, | 87 | + experimental_config: Optional[_ExperimentalConfig] = None, |
| 88 | + custom_trace_id_callback: Optional[Callable[[], str]] = None, | ||
| 85 | ) -> None: | 89 | ) -> None: |
| 86 | self._is_env_valid = check_msprof_env() | 90 | self._is_env_valid = check_msprof_env() |
| 87 | self.prof_path = "" | 91 | self.prof_path = "" |
| @@ -100,6 +104,8 @@ class _ProfInterface: | |||
| 100 | self.experimental_config = experimental_config | 104 | self.experimental_config = experimental_config |
| 101 | self.schedule = schedule | 105 | self.schedule = schedule |
| 102 | self.metadata = metadata | 106 | self.metadata = metadata |
| 107 | + self.custom_trace_id_callback = custom_trace_id_callback | ||
| 108 | + self.trace_id = "" | ||
| 103 | self.gc_detector = None | 109 | self.gc_detector = None |
| 104 | self._check_params() | 110 | self._check_params() |
| 105 | 111 | ||
| @@ -112,6 +118,7 @@ class _ProfInterface: | |||
| 112 | ProfPathCreator().create_prof_dir() | 118 | ProfPathCreator().create_prof_dir() |
| 113 | self.prof_path = ProfPathCreator().get_prof_dir() | 119 | self.prof_path = ProfPathCreator().get_prof_dir() |
| 114 | _init_profiler(self.prof_path, self.activities) | 120 | _init_profiler(self.prof_path, self.activities) |
| 121 | + self.trace_id = self.create_trace_id() | ||
| 115 | 122 | ||
| 116 | def warmup_trace(self): | 123 | def warmup_trace(self): |
| 117 | if not self._is_env_valid: | 124 | if not self._is_env_valid: |
| @@ -201,6 +208,32 @@ class _ProfInterface: | |||
| 201 | self.gc_detector.stop() | 208 | self.gc_detector.stop() |
| 202 | self.gc_detector = None | 209 | self.gc_detector = None |
| 203 | 210 | ||
| 211 | + def default_trace_id(self): | ||
| 212 | + # Generate a UUID | ||
| 213 | + uuid_raw = uuid.uuid4() | ||
| 214 | + return f"{uuid_raw.int:032X}" | ||
| 215 | + | ||
| 216 | + def create_trace_id(self): | ||
| 217 | + if not self.custom_trace_id_callback: | ||
| 218 | + return self.default_trace_id() | ||
| 219 | + if not isinstance(self.custom_trace_id_callback, Callable): | ||
| 220 | + print_warn_msg( | ||
| 221 | + "Parameter custom_trace_id_callback is not callable, reset it to default." | ||
| 222 | + ) | ||
| 223 | + return self.default_trace_id() | ||
| 224 | + try: | ||
| 225 | + trace_id = self.custom_trace_id_callback() | ||
| 226 | + if isinstance(trace_id, str) and len(trace_id) <= self.MAX_TRACE_ID_LEN: | ||
| 227 | + return trace_id | ||
| 228 | + print_warn_msg( | ||
| 229 | + f"Parameter custom_trace_id_callback should return str(max length: {self.MAX_TRACE_ID_LEN}), reset it to default." | ||
| 230 | + ) | ||
| 231 | + except Exception as e: | ||
| 232 | + print_warn_msg( | ||
| 233 | + f"Parameter custom_trace_id_callback raised an exception: {e}, reset it to default." | ||
| 234 | + ) | ||
| 235 | + return self.default_trace_id() | ||
| 236 | + | ||
⚠️ custom_trace_id_callback 返回非字符串值或抛异常时缺乏处理（同一段代码另有1个问题）
参考依据
xfeng 4月24日 评论： 4月24日 评论： | |||
| 204 | def _check_params(self): | 237 | def _check_params(self): |
| 205 | for activity in self.activities: | 238 | for activity in self.activities: |
| 206 | if activity in supported_activities(): | 239 | if activity in supported_activities(): |
| @@ -319,6 +352,7 @@ class _ProfInterface: | |||
| 319 | if Constant.Text in self.experimental_config.export_type: | 352 | if Constant.Text in self.experimental_config.export_type: |
| 320 | self.metadata.update(collect_env_vars()) | 353 | self.metadata.update(collect_env_vars()) |
| 321 | self._add_group_info_to_metadata() | 354 | self._add_group_info_to_metadata() |
| 355 | + self._add_trace_id_to_metadata() | ||
| 322 | if not self.metadata: | 356 | if not self.metadata: |
| 323 | return | 357 | return |
| 324 | if not ProfPathCreator().is_prof_inited: | 358 | if not ProfPathCreator().is_prof_inited: |
| @@ -372,6 +406,9 @@ class _ProfInterface: | |||
| 372 | except Exception as err: | 406 | except Exception as err: |
| 373 | print_warn_msg(f"Failed to get parallel group info, Exception: {str(err)}.") | 407 | print_warn_msg(f"Failed to get parallel group info, Exception: {str(err)}.") |
| 374 | 408 | ||
| 409 | + def _add_trace_id_to_metadata(self): | ||
| 410 | + self.metadata.update({self.TRACE_ID_KEY: self.trace_id}) | ||
| 411 | + | ||
| 375 | 412 | ||
| 376 | 413 | ||
| 377 | def supported_activities(): | 414 | def supported_activities(): |
| @@ -1,12 +1,11 @@ | |||
| 1 | from enum import Enum | 1 | from enum import Enum |
| 2 | 2 | ||
| 3 | from torch_npu.utils._error_code import ErrCode, prof_error | 3 | from torch_npu.utils._error_code import ErrCode, prof_error |
| 4 | + | ||
| 4 | from .analysis.prof_common_func._constant import print_warn_msg | 5 | from .analysis.prof_common_func._constant import print_warn_msg |
| 5 | 6 | ||
| 6 | -__all__ = [ | 7 | + |
| 7 | - 'ProfilerAction', | 8 | +__all__ = ["ProfilerAction", "Schedule"] |
| 8 | - 'Schedule' | ||
| 9 | -] | ||
| 10 | 9 | ||
| 11 | 10 | ||
| 12 | class ProfilerAction(Enum): | 11 | class ProfilerAction(Enum): |
| @@ -23,22 +22,44 @@ class Schedule: | |||
| 23 | ``active`` steps and then repeat the cycle starting with ``wait`` steps. The optional number | 22 | ``active`` steps and then repeat the cycle starting with ``wait`` steps. The optional number |
| 24 | of cycles is specified with the ``repeat`` parameter, the zero value means that | 23 | of cycles is specified with the ``repeat`` parameter, the zero value means that |
| 25 | the cycles will continue until the profiling is finished. | 24 | the cycles will continue until the profiling is finished. |
| 25 | + | ||
| 26 | + The ``skip_first_wait`` parameter controls whether the first ``wait`` stage should be skipped. | ||
| 27 | + This can be useful if a user wants to wait longer than ``skip_first`` between cycles, but not | ||
| 28 | + for the first profile. For example, if ``skip_first`` is 10 and ``wait`` is 20, the first cycle will | ||
| 29 | + wait 10 + 20 = 30 steps before warmup if ``skip_first_wait`` is zero, but will wait only 10 | ||
| 30 | + steps if ``skip_first_wait`` is non-zero. All subsequent cycles will then wait 20 steps between the | ||
| 31 | + last active and warmup. | ||
| 26 | """ | 32 | """ |
| 27 | - def __init__(self, wait: int, active: int, warmup: int = 0, repeat: int = 0, skip_first: int = 0) -> None: | 33 | + |
| 34 | + def __init__( | ||
| 35 | + self, | ||
| 36 | + wait: int, | ||
| 37 | + active: int, | ||
| 38 | + warmup: int = 0, | ||
| 39 | + repeat: int = 0, | ||
| 40 | + skip_first: int = 0, | ||
| 41 | + skip_first_wait: int = 0, | ||
| 42 | + ) -> None: | ||
| 28 | self.wait = wait | 43 | self.wait = wait |
| 29 | self.active = active | 44 | self.active = active |
| 30 | self.warmup = warmup | 45 | self.warmup = warmup |
| 31 | self.repeat = repeat | 46 | self.repeat = repeat |
| 32 | self.skip_first = skip_first | 47 | self.skip_first = skip_first |
| 48 | + self.skip_first_wait = skip_first_wait | ||
| 33 | self._check_params() | 49 | self._check_params() |
| 34 | 50 | ||
| 35 | def __call__(self, step: int) -> ProfilerAction: | 51 | def __call__(self, step: int) -> ProfilerAction: |
| 36 | if step < 0: | 52 | if step < 0: |
| 37 | - raise ValueError("Invalid parameter step, which must be not less than 0." + prof_error(ErrCode.VALUE)) | 53 | + raise ValueError( |
| 54 | + "Invalid parameter step, which must be not less than 0." | ||
| 55 | + + prof_error(ErrCode.VALUE) | ||
| 56 | + ) | ||
| 38 | if step < self.skip_first: | 57 | if step < self.skip_first: |
| 39 | return ProfilerAction.NONE | 58 | return ProfilerAction.NONE |
| 40 | else: | 59 | else: |
| 41 | step -= self.skip_first | 60 | step -= self.skip_first |
| 61 | + if self.skip_first_wait != 0: | ||
| 62 | + step += self.wait | ||
| 42 | num_steps = self.wait + self.warmup + self.active | 63 | num_steps = self.wait + self.warmup + self.active |
| 43 | if self.repeat > 0 and step / num_steps >= self.repeat: | 64 | if self.repeat > 0 and step / num_steps >= self.repeat: |
| 44 | return ProfilerAction.NONE | 65 | return ProfilerAction.NONE |
| @@ -70,9 +91,14 @@ class Schedule: | |||
| 70 | if not isinstance(self.skip_first, int) or self.skip_first < 0: | 91 | if not isinstance(self.skip_first, int) or self.skip_first < 0: |
| 71 | print_warn_msg("Invalid parameter skip_first, reset it to 0.") | 92 | print_warn_msg("Invalid parameter skip_first, reset it to 0.") |
| 72 | self.skip_first = 0 | 93 | self.skip_first = 0 |
| 94 | + if not isinstance(self.skip_first_wait, int): | ||
| 95 | + print_warn_msg("Invalid parameter skip_first_wait, reset it to 0.") | ||
| 96 | + self.skip_first_wait = 0 | ||
| 73 | 97 | ||
| 74 | if self.warmup == 0: | 98 | if self.warmup == 0: |
| 75 | - print_warn_msg("Profiler won't be using warmup, this can skew profiler results") | 99 | + print_warn_msg( |
| 100 | + "Profiler won't be using warmup, this can skew profiler results" | ||
| 101 | + ) | ||
| 76 | 102 | ||
| 77 | 103 | ||
| 78 | def _default_schedule_fn(_: int) -> ProfilerAction: | 104 | def _default_schedule_fn(_: int) -> ProfilerAction: |
⚠️ 测试覆盖不足：callback 异常场景未覆盖
测试仅覆盖了正常设置 callback 并获取 trace_id 的场景，未覆盖：(1) callback 抛异常时是否优雅降级；(2) callback 返回非字符串时是否降级；(3) 不设置 callback 时默认生成 UUID 的场景。边界条件测试缺失可能导致生产环境问题。
Request Changes
参考依据