mindspore 2.4.1__cp39-none-any.whl → 2.5.0__cp39-none-any.whl

mindspore/profiler/analysis/viewer/ascend_op_memory_viewer.py

@@ -0,0 +1,322 @@
+# Copyright 2024 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.
+# ============================================================================
+"""Ascend op memory viewer"""
+import os
+import struct
+from collections import defaultdict
+from enum import Enum
+from typing import Any, Dict, List
+from abc import ABC
+
+from mindspore.profiler.parser.ascend_analysis.tlv_decoder import TLVDecoder
+from mindspore.profiler.common.file_manager import FileManager
+from mindspore.profiler.common.log import ProfilerLogger
+
+
+class OpMemoryIndexEnum(Enum):
+    """Op memory index defining."""
+
+    DEVICE_ID = 0
+    TID = 1
+    PID = 2
+    CREATE_AT = 3
+    ADDR = 4
+    SIZE = 5
+    USED_SIZE = 6
+    PEAK_SIZE = 7
+    ALLOC_SIZE = 8
+    USED_BY_EVENT_SZIE = 9
+    EAGER_FREE_SIZE = 10
+    STEAM_PTR = 11
+    STREAM_ID = 12
+    FROM_PERSISTENT = 13
+    IS_PERSISTENT = 14
+    RUN_MODE = 15
+    ALLOC_TYPE = 16
+    OWNER = 17
+
+
+class BaseEvent(ABC):
+    """Base class for all event types."""
+
+    def __init__(self, data: Dict):
+        if not isinstance(data, dict):
+            raise TypeError("Input data must be dict.")
+        self._origin_data = data
+
+
+class OpMemoryEvent(BaseEvent):
+    """Op memory event."""
+
+    FIX_DATA_FORMAT = "<i11QI4B"
+    FIX_DATA_SIZE = struct.calcsize(FIX_DATA_FORMAT)
+    FREE_VALUE = 18446744073709551615  # 2^64 - 1
+    NAME_KEY = 3
+
+    def __init__(self, data: Dict):
+        super().__init__(data)
+        self.fix_size_data = struct.unpack(
+            self.FIX_DATA_FORMAT, self._origin_data.get("fix_size_bytes")
+        )
+
+    @property
+    def device_id(self):
+        """Get device id."""
+        return self.fix_size_data[OpMemoryIndexEnum.DEVICE_ID.value]
+
+    @property
+    def tid(self):
+        """Get tid."""
+        return self.fix_size_data[OpMemoryIndexEnum.TID.value]
+
+    @property
+    def pid(self):
+        """Get pid."""
+        return self.fix_size_data[OpMemoryIndexEnum.PID.value]
+
+    @property
+    def create_at(self):
+        """Get create at."""
+        return self.fix_size_data[OpMemoryIndexEnum.CREATE_AT.value]
+
+    @property
+    def addr(self):
+        """Get addr."""
+        return self.fix_size_data[OpMemoryIndexEnum.ADDR.value]
+
+    @property
+    def size(self):
+        """Get size."""
+        return self.fix_size_data[OpMemoryIndexEnum.SIZE.value]
+
+    @property
+    def used_size(self):
+        """Get used size."""
+        return self.fix_size_data[OpMemoryIndexEnum.USED_SIZE.value]
+
+    @property
+    def peak_size(self):
+        """Get peak size."""
+        return self.fix_size_data[OpMemoryIndexEnum.PEAK_SIZE.value]
+
+    @property
+    def alloc_size(self):
+        """Get alloc size."""
+        return self.fix_size_data[OpMemoryIndexEnum.ALLOC_SIZE.value]
+
+    @property
+    def used_by_event_size(self):
+        """Get used by event size."""
+        return self.fix_size_data[OpMemoryIndexEnum.USED_BY_EVENT_SZIE.value]
+
+    @property
+    def eager_free_size(self):
+        """Get eager free size."""
+        return self.fix_size_data[OpMemoryIndexEnum.EAGER_FREE_SIZE.value]
+
+    @property
+    def stream_ptr(self):
+        """Get stream ptr."""
+        return self.fix_size_data[OpMemoryIndexEnum.STEAM_PTR.value]
+
+    @property
+    def stream_id(self):
+        """Get stream id."""
+        return self.fix_size_data[OpMemoryIndexEnum.STREAM_ID.value]
+
+    @property
+    def from_persistent(self):
+        """Get from persistent."""
+        return self.fix_size_data[OpMemoryIndexEnum.FROM_PERSISTENT.value]
+
+    @property
+    def is_persistent(self):
+        """Get is persistent."""
+        return self.fix_size_data[OpMemoryIndexEnum.IS_PERSISTENT.value]
+
+    @property
+    def run_mode(self):
+        """Get run mode."""
+        return self.fix_size_data[OpMemoryIndexEnum.RUN_MODE.value]
+
+    @property
+    def alloc_type(self):
+        """Get alloc type."""
+        return self.fix_size_data[OpMemoryIndexEnum.ALLOC_TYPE.value]
+
+    @property
+    def owner(self):
+        """Get owner."""
+        return self._origin_data.get(self.NAME_KEY, "")
+
+    @property
+    def is_alloc(self):
+        """Get is alloc."""
+        return self.size != self.FREE_VALUE
+
+
+class AscendOpMemoryViewer:
+    """
+    Ascend op memory viewer.
+    """
+
+    FWK_BINARY_FILE_NAME = "mindspore.memory_usage"
+    OUTPUT_FILE_NAME = "operator_memory.csv"
+    HEADERS = [
+        "Name",
+        "Size(KB)",
+        "Allocation Time(us)",
+        "Release Time(us)",
+        "Active Release Time(us)",
+        "Duration(us)",
+        "Active Duration(us)",
+        "Allocation Total Allocated(MB)",
+        "Allocation Total Reserved(MB)",
+        "Allocation Total Active(MB)",
+        "Release Total Allocated(MB)",
+        "Release Total Reserved(MB)",
+        "Release Total Active(MB)",
+        "Stream Ptr",
+        "Device Type",
+    ]
+    DEVICE_TYPE_FMT = "NPU:{}"
+    NS_TO_US = 1024
+    BYTES_TO_KB = 1024
+    BYTES_TO_MB = 1024 * 1024
+    EMPTY_VALUE = "N/A"
+    ALLOC_TIME_INDEX = 2
+
+    def __init__(self, **kwargs):
+        self._enable_profile_memory = kwargs.get("profile_memory", False)
+        self._framework_path = kwargs.get("framework_path")
+        self._ascend_profiler_output_path = kwargs.get("ascend_profiler_output_path")
+        self._ascend_ms_dir = kwargs.get("ascend_ms_dir")
+        ProfilerLogger.init(self._ascend_ms_dir)
+        self._logger = ProfilerLogger.get_instance()
+        self._op_memory_events = None
+        self._op_memory_data = []
+        self._addr_to_event_map = defaultdict(list)
+
+    def save(self, data: Any = None):
+        """
+        Save step trace time data to csv file
+        """
+        self._logger.info("AscendOpMemoryViewer start")
+        if not self._enable_profile_memory:
+            return
+
+        try:
+            self._read_fwk_binary_file()
+            self._calculate_op_memory_data()
+            self._write_data()
+        except Exception as e:  # pylint: disable=W0703
+            self._logger.error("Failed to save op memory data: %s", str(e), exc_info=True)
+        self._logger.info("AscendOpMemoryViewer end")
+
+    def _read_fwk_binary_file(self):
+        """
+        Read fwk binary file
+        """
+        self._logger.info("Read fwk binary file start")
+        op_name_file_path = os.path.join(self._framework_path, self.FWK_BINARY_FILE_NAME)
+        raw_bin_data = FileManager.read_file_content(op_name_file_path, mode="rb")
+        self._op_memory_events = TLVDecoder.decode(
+            raw_bin_data, OpMemoryEvent, OpMemoryEvent.FIX_DATA_SIZE
+        )
+        self._op_memory_events = sorted(self._op_memory_events, key=lambda x: x.create_at)
+        self._logger.info("Read fwk binary file done, %d events", len(self._op_memory_events))
+
+    def _calculate_op_memory_data(self):
+        """
+        Calculate op memory data
+        """
+        self._logger.info("Calculate op memory data start")
+        if not self._op_memory_events:
+            self._logger.info("No op memory events")
+            return
+
+        for event in self._op_memory_events:
+            self._addr_to_event_map[event.addr].append(event)
+
+        for _, event_list in self._addr_to_event_map.items():
+            row_data_list = self._get_op_mem_row_data(event_list)
+            self._op_memory_data.extend(row_data_list)
+
+        self._op_memory_data = sorted(self._op_memory_data, key=lambda x: x[self.ALLOC_TIME_INDEX])
+        self._logger.info("Calculate op memory data done")
+
+    def _combine_alloc_and_free_event(self, alloc_event: OpMemoryEvent, free_event=None):
+        """
+        Combine alloc and free event
+        """
+        if not alloc_event:
+            self._logger.error("Alloc event is None")
+            return []
+
+        return [
+            alloc_event.owner, # "Name"
+            alloc_event.size / self.BYTES_TO_KB, # "Size(KB)"
+            alloc_event.create_at / self.NS_TO_US, # "Allocation Time(us)"
+            self.EMPTY_VALUE if free_event is None else free_event.create_at / self.NS_TO_US, # "Release Time(us)"
+            self.EMPTY_VALUE, # "Active Release Time(us)"
+            self.EMPTY_VALUE if free_event is None else (free_event.create_at - alloc_event.create_at) / self.NS_TO_US, # "Duration(us)"
+            self.EMPTY_VALUE, # "Active Duration(us)"
+            alloc_event.alloc_size / self.BYTES_TO_MB, # "Allocation Total Allocated(MB)"
+            alloc_event.used_size / self.BYTES_TO_MB, # "Allocation Total Reserved(MB)"
+            alloc_event.alloc_size / self.BYTES_TO_MB, # "Allocation Total Active(MB)"
+            self.EMPTY_VALUE if free_event is None else free_event.alloc_size / self.BYTES_TO_MB, # "Release Total Allocated(MB)"
+            self.EMPTY_VALUE if free_event is None else free_event.used_size / self.BYTES_TO_MB, # "Release Total Reserved(MB)"
+            self.EMPTY_VALUE if free_event is None else free_event.alloc_size / self.BYTES_TO_MB, # "Release Total Active(MB)"
+            alloc_event.stream_ptr, # "Stream Ptr"
+            self.DEVICE_TYPE_FMT.format(alloc_event.device_id), # "Device Type"
+        ]
+
+    def _get_op_mem_row_data(self, event_list: List[OpMemoryEvent]):
+        """
+        Get op memory row data
+        """
+        res = []
+
+        if not event_list:
+            self._logger.error("Event list length is less than 1")
+            return res
+
+        start_index = 0 if event_list[0].is_alloc else 1
+        alloc_event, free_event = None, None
+        for event in event_list[start_index:]:
+            if event.is_alloc:
+                alloc_event = event
+            else:
+                free_event = event
+
+            if alloc_event and free_event:
+                res.append(self._combine_alloc_and_free_event(alloc_event, free_event))
+                alloc_event, free_event = None, None
+            elif alloc_event is None and free_event:
+                self._logger.error("Alloc event is None, but free event is not None")
+
+        if alloc_event:
+            res.append(self._combine_alloc_and_free_event(alloc_event))
+
+        return res
+
+    def _write_data(self):
+        """
+        Write data to csv file
+        """
+        self._logger.info("Write data to csv file start")
+        save_path = os.path.join(self._ascend_profiler_output_path, self.OUTPUT_FILE_NAME)
+        FileManager.create_csv_file(save_path, self._op_memory_data, self.HEADERS)
+        self._logger.info("Write data to csv file done, %d rows, save path: %s", len(self._op_memory_data), save_path)

mindspore/profiler/analysis/viewer/ascend_step_trace_time_viewer.py

@@ -0,0 +1,265 @@
+# Copyright 2024 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.
+# ============================================================================
+"""Ascend Step Trace Time Viewer"""
+import os
+import re
+from decimal import Decimal
+from enum import Enum
+from typing import List, Any, Tuple
+
+import numpy as np
+
+from mindspore import log as logger
+from mindspore.profiler.analysis.viewer.base_viewer import BaseViewer
+from mindspore.profiler.common.file_manager import FileManager
+from mindspore.profiler.common.constant import TimelineLayerName, OverlapAnalysisTidName, ProfilerLevel
+from mindspore.profiler.analysis.parser.timeline_event.msprof_event import (
+    MsprofCompleteEvent,
+)
+from mindspore.profiler.analysis.parser.timeline_event.timeline_event_pool import (
+    TimelineEventPool,
+)
+from mindspore.profiler.analysis.parser.timeline_assembly_factory.trace_view_container import (
+    TraceViewContainer,
+)
+from mindspore.profiler.common.log import ProfilerLogger
+
+
+class StepTraceTimeHeaders(Enum):
+    """Step trace time headers"""
+    STEP = "Step"
+    COMPUTING = "Computing"
+    COMMUNICATION_NOT_OVERLAPPED = "Communication(Not Overlapped)"
+    OVERLAPPED = "Overlapped"
+    COMMUNICATION = "Communication"
+    FREE = "Free"
+    STAGE = "Stage"
+    BUBBLE = "Bubble"
+    COMMUNICATION_NOT_OVERLAPPED_EXCLUDE_RECEIVE = "Communication(Not Overlapped and Exclude Receive)"
+    PREPARING = "Preparing"
+
+
+class AscendStepTraceTimeViewer(BaseViewer):
+    """Ascend Step Trace Time Viewer"""
+
+    STEP_TRACE_TIME_FILE_NAME = "step_trace_time.csv"
+    STEP_TRACE_TIME_HEADERS = [header.value for header in StepTraceTimeHeaders]
+
+    # HCCL Send, Recv op pattern
+    PP_OP_PATTERN = (
+        # eg: hcom_BatchSendRecv__101_0_1
+        re.compile(r"hcom_\w+SendRecv__\d+"),
+        # eg: hcom_send__101_0_1
+        re.compile(r"hcom_send__\d+"),
+        # eg: hcom_receive__101_0_1
+        re.compile(r"hcom_receive__\d+"),
+        re.compile(r"Receive-op"),
+        re.compile(r"Send-op"),
+    )
+
+    # numpy array dtype
+    OVERLAP_DTYPE = np.dtype([("ts", object), ("dur", object)])
+    HCCL_DTYPE = np.dtype([("name", object), ("ts", object), ("dur", object)])
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self._save_path = os.path.join(
+            kwargs.get("ascend_profiler_output_path"), self.STEP_TRACE_TIME_FILE_NAME
+        )
+        self._profiler_level = kwargs.get("profiler_level")
+        self._ascend_ms_dir = kwargs.get("ascend_ms_dir")
+        ProfilerLogger.init(self._ascend_ms_dir)
+        self._logger = ProfilerLogger.get_instance()
+        self.step_trace_time_data = {}
+        self.trace_container: TraceViewContainer = None
+        self.hccl_pool: TimelineEventPool = None
+        self.overlap_pool: TimelineEventPool = None
+        # HCCL events
+        self.hccl_events: List[MsprofCompleteEvent] = None
+        # Overlap analysis events
+        self.computing_events: List[MsprofCompleteEvent] = None
+        self.communication_events: List[MsprofCompleteEvent] = None
+        self.communication_not_overlapped_events: List[MsprofCompleteEvent] = None
+        self.free_events: List[MsprofCompleteEvent] = None
+        # Overlap analysis numpy array
+        self.computing_np: np.ndarray = None
+        self.communication_np: np.ndarray = None
+        self.communication_not_overlapped_np: np.ndarray = None
+        self.free_np: np.ndarray = None
+        # HCCL numpy array
+        self.hccl_events_np: np.ndarray = None
+
+    def save(self, data: Any):
+        """
+        Save step trace time data to csv file
+        """
+        self._logger.info("AscendStepTraceTimeViewer start")
+        if self._profiler_level == ProfilerLevel.LevelNone.value:
+            return
+        try:
+            self._check_input_data(data)
+            self._convert_events_to_numpy()
+            self._calculate_step_trace_time()
+            self._write_data()
+        except Exception as e:  # pylint: disable=W0703
+            self._logger.error("Failed to save step trace time data, error: %s", str(e), exc_info=True)
+        self._logger.info("AscendStepTraceTimeViewer end")
+
+    def _write_data(self):
+        """
+        Write step trace time data to csv file
+        """
+        self._logger.info("Write step trace time data start")
+        data = [[str(self.step_trace_time_data.get(header, "")) for header in self.STEP_TRACE_TIME_HEADERS]]
+        FileManager.create_csv_file(
+            self._save_path,
+            data,
+            self.STEP_TRACE_TIME_HEADERS,
+        )
+        self._logger.info("Write step trace time data done, %d rows saved, save path: %s", len(data), self._save_path)
+
+    def _check_input_data(self, data: Any):
+        """
+        Check input data and initialize data
+        """
+        self.trace_container: TraceViewContainer = data.get(
+            "trace_view_container", None
+        )
+
+        if self.trace_container is None:
+            raise ValueError("trace is empty")
+
+        self.overlap_pool: TimelineEventPool = self.trace_container.get_pool_by_name(
+            TimelineLayerName.OVERLAP_ANALYSIS.value
+        )
+        self.hccl_pool: TimelineEventPool = self.trace_container.get_pool_by_name(
+            TimelineLayerName.HCCL.value
+        )
+
+        if self.overlap_pool is None:
+            raise ValueError("overlap pool is empty")
+
+        self.computing_events: List[MsprofCompleteEvent] = (
+            self.overlap_pool.get_events_by_name(OverlapAnalysisTidName.COMPUTING.value)
+        )
+        self.communication_events: List[MsprofCompleteEvent] = (
+            self.overlap_pool.get_events_by_name(
+                OverlapAnalysisTidName.COMMUNICATION.value
+            )
+        )
+        self.communication_not_overlapped_events: List[MsprofCompleteEvent] = (
+            self.overlap_pool.get_events_by_name(
+                OverlapAnalysisTidName.COMMUNICATION_NOT_OVERLAP.value
+            )
+        )
+        self.free_events: List[MsprofCompleteEvent] = (
+            self.overlap_pool.get_events_by_name(OverlapAnalysisTidName.FREE.value)
+        )
+        if self.hccl_pool is not None:
+            self.hccl_events: List[MsprofCompleteEvent] = (
+                self.hccl_pool.get_complete_events()
+            )
+
+    def _convert_overlap_events_to_numpy(
+            self, events: List[MsprofCompleteEvent], dtype
+    ):
+        """
+        Convert overlap events to numpy array
+        """
+        return np.array([(event.ts, event.dur) for event in events], dtype=dtype)
+
+    def _convert_events_to_numpy(self):
+        """
+        Convert events to numpy array
+        """
+        self.computing_np = self._convert_overlap_events_to_numpy(
+            self.computing_events, self.OVERLAP_DTYPE
+        )
+        self.communication_np = self._convert_overlap_events_to_numpy(
+            self.communication_events, self.OVERLAP_DTYPE
+        )
+        self.communication_not_overlapped_np = self._convert_overlap_events_to_numpy(
+            self.communication_not_overlapped_events, self.OVERLAP_DTYPE
+        )
+        self.free_np = self._convert_overlap_events_to_numpy(
+            self.free_events, self.OVERLAP_DTYPE
+        )
+        self.computing_np = np.sort(self.computing_np, order="ts")
+        self.communication_np = np.sort(self.communication_np, order="ts")
+        self.communication_not_overlapped_np = np.sort(
+            self.communication_not_overlapped_np, order="ts"
+        )
+        self.free_np = np.sort(self.free_np, order="ts")
+
+        if self.hccl_events is not None:
+            self.hccl_events_np = np.array(
+                [(event.name, event.ts, event.dur) for event in self.hccl_events],
+                dtype=self.HCCL_DTYPE,
+            )
+            self.hccl_events_np = np.sort(self.hccl_events_np, order="ts")
+
+    def _calculate_step_trace_time(self):
+        """
+        Calculate step trace time data
+        """
+        self.step_trace_time_data[StepTraceTimeHeaders.STEP.value] = "0"
+        self.step_trace_time_data[StepTraceTimeHeaders.COMPUTING.value] = np.sum(self.computing_np["dur"])
+        self.step_trace_time_data[StepTraceTimeHeaders.COMMUNICATION.value] = np.sum(self.communication_np["dur"])
+        self.step_trace_time_data[StepTraceTimeHeaders.COMMUNICATION_NOT_OVERLAPPED.value] = np.sum(
+            self.communication_not_overlapped_np["dur"]
+        )
+        self.step_trace_time_data[StepTraceTimeHeaders.OVERLAPPED.value] = (
+            self.step_trace_time_data[StepTraceTimeHeaders.COMMUNICATION.value]
+            - self.step_trace_time_data[StepTraceTimeHeaders.COMMUNICATION_NOT_OVERLAPPED.value]
+        )
+        self.step_trace_time_data[StepTraceTimeHeaders.FREE.value] = np.sum(self.free_np["dur"])
+        (
+            self.step_trace_time_data[StepTraceTimeHeaders.STAGE.value],
+            self.step_trace_time_data[StepTraceTimeHeaders.BUBBLE.value],
+        ) = self._calculate_stage_bubble()
+
+        self.step_trace_time_data[StepTraceTimeHeaders.COMMUNICATION_NOT_OVERLAPPED_EXCLUDE_RECEIVE.value] = (
+            self.step_trace_time_data[StepTraceTimeHeaders.COMMUNICATION_NOT_OVERLAPPED.value]
+            - self.step_trace_time_data[StepTraceTimeHeaders.BUBBLE.value]
+        )
+        self.step_trace_time_data[StepTraceTimeHeaders.PREPARING.value] = "0"
+
+    def _calculate_stage_bubble(self) -> Tuple[Decimal, Decimal]:
+        """
+        Calculate stage and bubble time
+        """
+        if self.hccl_events is None:
+            logger.info("HCCL events is empty, skip calculate stage and bubble")
+            return Decimal(0), Decimal(0)
+
+        total_hccl_time = self.hccl_events_np["ts"][-1] - self.hccl_events_np["ts"][0]
+        bubble_time = np.sum(
+            self.hccl_events_np["dur"][
+                np.array(
+                    [
+                        self._is_send_recv_op(name)
+                        for name in self.hccl_events_np["name"]
+                    ]
+                )
+            ]
+        )
+        stage_time = total_hccl_time - bubble_time
+        return stage_time, bubble_time
+
+    def _is_send_recv_op(self, op_name: str) -> bool:
+        """
+        Check if the op is a send or recv op
+        """
+        return any(pattern.match(op_name) for pattern in self.PP_OP_PATTERN)

mindspore/profiler/analysis/viewer/ascend_timeline_viewer.py

@@ -0,0 +1,58 @@
+# Copyright 2024 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.
+# ============================================================================
+"""ascend timeline viewer"""
+import os
+from typing import List, Dict
+
+from mindspore.profiler.analysis.viewer.base_viewer import BaseViewer
+from mindspore.profiler.common.file_manager import FileManager
+from mindspore.profiler.common.log import ProfilerLogger
+
+
+class AscendTimelineViewer(BaseViewer):
+    """Ascend Timeline Viewer for analyzing and saving timeline data."""
+
+    _TRACE_VIEW_FILE_NAME = 'trace_view.json'
+
+    def __init__(self, **kwargs):
+        """Initialize the AscendTimelineViewer."""
+        super().__init__()
+        self._save_path = os.path.join(
+            kwargs.get("ascend_profiler_output_path"),
+            self._TRACE_VIEW_FILE_NAME
+        )
+        self._ascend_ms_dir = kwargs.get("ascend_ms_dir")
+        ProfilerLogger.init(self._ascend_ms_dir)
+        self._logger = ProfilerLogger.get_instance()
+
+    def save(self, data: Dict) -> None:
+        """Get the input data and save the timeline data."""
+        self._logger.info("AscendTimelineViewer start")
+        try:
+            trace_view_container = data.get("trace_view_container", None)
+            if not trace_view_container:
+                raise RuntimeError("The trace view container is None, Failed to save trace_view.json.")
+            trace_view_data = trace_view_container.get_trace_view()
+            self._save_data(trace_view_data)
+            self._logger.info("Trace viewer save trace_view.json done")
+        except Exception as e: # pylint: disable=W0703
+            self._logger.error("Failed to save trace_view.json: %s", e, exc_info=True)
+        self._logger.info("AscendTimelineViewer end")
+
+    def _save_data(self, timeline_data: List[Dict]) -> None:
+        """Save the timeline data to a JSON file."""
+        self._logger.info("Trace view saved start")
+        FileManager.create_json_file(self._save_path, timeline_data)
+        self._logger.info("Trace view saved done, save path: %s", self._save_path)

mindspore/profiler/analysis/viewer/base_viewer.py

@@ -0,0 +1,26 @@
+# Copyright 2024 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.
+# ============================================================================
+"""BaseViewer class"""
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseViewer(ABC):
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def save(self, data: Any):
+        raise NotImplementedError("Subclasses should implement this!")