[Typing][B-79,B-80,B-81] Add type annotations for python/paddle/profiler/{profiler,profiler_statistic,utils}.py (#67391)

megemini · web-flow · commit 3a58f35d06e6 · 2024-08-15T11:58:47.000+08:00
* [Add] typing

* [Fix] typing self
diff --git a/python/paddle/profiler/profiler.py b/python/paddle/profiler/profiler.py
@@ -19,7 +19,7 @@
 import os
 import socket
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Callable, Literal
 from warnings import warn
 
 import paddle
@@ -45,6 +45,11 @@
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
+    from types import TracebackType
+
+    from typing_extensions import Self
+
+    from paddle.base.core import _ProfilerResult
 
 
 class SummaryView(Enum):
@@ -128,7 +133,7 @@ def make_scheduler(
     record: int,
     repeat: int = 0,
     skip_first: int = 0,
-) -> Callable:
+) -> Callable[[int], ProfilerState]:
     r"""
     Return a scheduler function, which scheduler the :ref:`ProfilerState <api_paddle_profiler_ProfilerState>` according to the setting.
     The state transform confirms to:
@@ -221,7 +226,7 @@ def _default_state_scheduler(step: int):
 
 def export_chrome_tracing(
     dir_name: str, worker_name: str | None = None
-) -> Callable:
+) -> Callable[[Profiler], None]:
     r"""
     Return a callable, used for outputing tracing data to chrome tracing format file.
     The output file will be saved in directory ``dir_name``, and file name will be set as `worker_name`.
@@ -272,7 +277,9 @@ def handle_fn(prof):
     return handle_fn
 
 
-def export_protobuf(dir_name: str, worker_name: str | None = None) -> Callable:
+def export_protobuf(
+    dir_name: str, worker_name: str | None = None
+) -> Callable[[Profiler], None]:
     r"""
     Return a callable, used for outputing tracing data to protobuf file.
     The output file will be saved in directory ``dir_name``, and file name will be set as ``worker_name``.
@@ -478,19 +485,36 @@ class Profiler:
                 >>> # |       ips       |    1086.42904   |    1227.30604   |    959.92796    |
     """
 
+    targets: Iterable[ProfilerTarget]
+    profiler: _Profiler
+    scheduler: Callable[[int], ProfilerState]
+    on_trace_ready: Callable[[Profiler], None]
+    step_num: int
+    previous_state: ProfilerState
+    current_state: ProfilerState
+    record_event: RecordEvent
+    profiler_result: _ProfilerResult
+    timer_only: bool
+    record_shapes: bool
+    profile_memory: bool
+    with_flops: bool
+    emit_nvtx: bool
+
     def __init__(
         self,
         *,
         targets: Iterable[ProfilerTarget] | None = None,
-        scheduler: Callable[[int], ProfilerState] | tuple | None = None,
-        on_trace_ready: Callable[..., Any] | None = None,
-        record_shapes: bool | None = False,
-        profile_memory: bool | None = False,
-        timer_only: bool | None = False,
-        emit_nvtx: bool | None = False,
-        custom_device_types: list | None = [],
-        with_flops: bool | None = False,
-    ):
+        scheduler: (
+            Callable[[int], ProfilerState] | tuple[int, int] | None
+        ) = None,
+        on_trace_ready: Callable[[Profiler], None] | None = None,
+        record_shapes: bool = False,
+        profile_memory: bool = False,
+        timer_only: bool = False,
+        emit_nvtx: bool = False,
+        custom_device_types: list[str] = [],
+        with_flops: bool = False,
+    ) -> None:
         supported_targets = _get_supported_targets()
         if targets:
             self.targets = set(targets)
@@ -553,14 +577,19 @@ def __init__(
         self.with_flops = with_flops
         self.emit_nvtx = emit_nvtx
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
         self.stop()
 
-    def start(self):
+    def start(self) -> None:
         r'''
         Start profiler and enter the first profiler step(0).
         State transformed from CLOSED to self.current_state and trigger corresponding action.
@@ -609,7 +638,7 @@ def start(self):
         )
         self.record_event.begin()
 
-    def stop(self):
+    def stop(self) -> None:
         r'''
         Stop profiler and State transformed from self.current_state to CLOSED.
         Trigger corresponding action and post-process profiler result using self.on_trace_ready if result exists.
@@ -659,7 +688,7 @@ def stop(self):
                 self.on_trace_ready(self)
         utils._is_profiler_used = False
 
-    def step(self, num_samples: int | None = None):
+    def step(self, num_samples: int | None = None) -> None:
         r"""
         Signals the profiler that the next profiling step has started.
         Get the new ProfilerState and trigger corresponding action.
@@ -703,7 +732,7 @@ def step(self, num_samples: int | None = None):
         )
         self.record_event.begin()
 
-    def step_info(self, unit=None):
+    def step_info(self, unit: str | None = None) -> str:
         r"""
         Get statistics for current step. If the function is called at certain iteration
         intervals, the result is the average of all steps between the previous call and
@@ -821,7 +850,7 @@ def _trigger_action(self):
             if self.on_trace_ready:
                 self.on_trace_ready(self)
 
-    def export(self, path="", format="json"):
+    def export(self, path: str = "", format: str = "json") -> None:
         r"""
         Exports the tracing data to file.
 
@@ -853,20 +882,20 @@ def export(self, path="", format="json"):
 
     def summary(
         self,
-        sorted_by=SortedKeys.CPUTotal,
-        op_detail=True,
-        thread_sep=False,
-        time_unit='ms',
-        views=None,
-    ):
+        sorted_by: SortedKeys = SortedKeys.CPUTotal,
+        op_detail: bool = True,
+        thread_sep: bool = False,
+        time_unit: Literal['s', 'ms', 'us', 'ns'] = 'ms',
+        views: SummaryView | list[SummaryView] | None = None,
+    ) -> None:
         r"""
         Print the Summary table. Currently support overview, model, distributed, operator, memory manipulation and user-defined summary.
 
         Args:
             sorted_by( :ref:`SortedKeys <api_paddle_profiler_SortedKeys>` , optional): how to rank the op table items, default value is SortedKeys.CPUTotal.
             op_detail(bool, optional): expand each operator detail information, default value is True.
             thread_sep(bool, optional): print op table each thread, default value is False.
-            time_unit(str, optional): time unit for display, can be chosen form ['s', 'ms', 'us', 'ns'], default value is 'ms'.
+            time_unit(str, optional): time unit for display, can be chosen from ['s', 'ms', 'us', 'ns'], default value is 'ms'.
             views(SummaryView|list[SummaryView], optional): summary tables to print, default to None means all views to be printed.
 
         Examples:
diff --git a/python/paddle/profiler/utils.py b/python/paddle/profiler/utils.py
@@ -26,6 +26,10 @@
 if TYPE_CHECKING:
     import types
 
+    from typing_extensions import Self
+
+    from paddle.base.core import _ProfilerResult
+
 _is_profiler_used = False
 _has_optimizer_wrapped = False
 
@@ -74,16 +78,20 @@ class RecordEvent(ContextDecorator):
         RecordEvent will take effect only when :ref:`Profiler <api_paddle_profiler_Profiler>` is on and at the state of `RECORD`.
     """
 
+    name: str
+    event_type: TracerEventType
+    event: _RecordEvent | None
+
     def __init__(
         self,
         name: str,
         event_type: TracerEventType = TracerEventType.PythonUserDefined,
-    ):
+    ) -> None:
         self.name = name
         self.event_type = event_type
         self.event = None
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         self.begin()
         return self
 
@@ -92,10 +100,10 @@ def __exit__(
         exc_type: type[BaseException] | None,
         exc_value: BaseException | None,
         traceback: types.TracebackType | None,
-    ):
+    ) -> None:
         self.end()
 
-    def begin(self):
+    def begin(self) -> None:
         r"""
         Record the time of beginning.
 
@@ -127,7 +135,7 @@ def begin(self):
         else:
             self.event = _RecordEvent(self.name, self.event_type)
 
-    def end(self):
+    def end(self) -> None:
         r"""
         Record the time of ending.
 
@@ -150,7 +158,7 @@ def end(self):
             self.event.end()
 
 
-def load_profiler_result(filename: str):
+def load_profiler_result(filename: str) -> _ProfilerResult:
     r"""
     Load dumped profiler data back to memory.
 
