[Typing][C-107] Add type annotations for python/paddle/base/compiler.py (#67767)

megemini · web-flow · commit 80c8760ae85f · 2024-08-29T20:27:31.000+08:00
diff --git a/python/paddle/base/compiler.py b/python/paddle/base/compiler.py
@@ -12,11 +12,57 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import sys
+from typing import TYPE_CHECKING, Any, TypedDict
 
 from . import core, framework
 from .framework import cpu_places, cuda_places, xpu_places
 
+if TYPE_CHECKING:
+    from paddle.base.core import Graph, _Scope
+    from paddle.optimizer import Optimizer
+    from paddle.static import Program
+
+    class _CustomOp(TypedDict):
+        paddle_op: str
+        popart_op: str
+        domain: str
+        version: int
+
+    class _IpuStrategyOptions(TypedDict, total=False):
+        is_training: bool
+        need_avg_shard: bool
+        enable_fp16: bool
+        use_no_bias_optimizer: bool
+        enable_distribution: bool
+        scaled_optimizer_state: bool
+        is_dynamic: bool
+        enable_model_runtime_executor: bool
+        num_ipus: int
+        batches_per_step: int
+        micro_batch_size: int
+        random_seed: int
+        tiles_per_ipu: int
+        num_buffers: int
+        available_memory_proportion: float
+        loss_scaling: float
+        max_weight_norm: float
+        timeout_ms: float
+        lr: float
+        accl1_type: str
+        accl2_type: str
+        accl3_type: str
+        onnx_dump_path: str
+        weight_decay_mode: str
+        enable_pipelining: bool
+        enable_gradient_accumulation: bool
+        accumulation_factor: int
+        enable_manual_shard: bool
+        custom_op: _CustomOp
+
+
 __all__ = []
 
 BuildStrategy = core.CompiledProgram.BuildStrategy
@@ -124,7 +170,11 @@ class CompiledProgram:
             ...                     fetch_list=[loss.name])
     """
 
-    def __init__(self, program_or_graph, build_strategy=None):
+    def __init__(
+        self,
+        program_or_graph: Graph | Program,
+        build_strategy: BuildStrategy | None = None,
+    ) -> None:
         if isinstance(program_or_graph, core.Graph):
             self._graph = program_or_graph
             # don't not create a new program here.
@@ -554,7 +604,11 @@ class IpuStrategy:
             >>> ipu_strategy = static.IpuStrategy()
     """
 
-    def __init__(self):
+    has_custom_ops: bool
+    custom_op_names: list[str]
+    need_compile: bool
+
+    def __init__(self) -> None:
         if core.is_compiled_with_ipu():
             self._ipu_strategy = core.IpuStrategy()
             default_options = {
@@ -578,7 +632,7 @@ def __init__(self):
         if in_dynamic_mode():
             self.register_patch()
 
-    def register_patch(self):
+    def register_patch(self) -> None:
         """
         Register patch function to support dynamic to static on IPU. This operation would break the dy2static functionality on CPU.
         Use `release_patch` to release the patch.
@@ -597,7 +651,7 @@ def register_patch(self):
         """
         IpuDynamicPatcher.register_patch(self)
 
-    def release_patch(self):
+    def release_patch(self) -> None:
         """
         Release the registered IPU functions.
 
@@ -615,7 +669,7 @@ def release_patch(self):
         """
         IpuDynamicPatcher.release_patch()
 
-    def set_optimizer(self, optimizer):
+    def set_optimizer(self, optimizer: Optimizer) -> None:
         """
         Set optimizer to ipu_strategy in dynamic mode.
 
@@ -647,7 +701,7 @@ def set_optimizer(self, optimizer):
         else:
             raise RuntimeError("Only needs to set optimizer in dynamic mode.")
 
-    def parse_optimizer(self, optimizer):
+    def parse_optimizer(self, optimizer: Optimizer) -> _IpuStrategyOptions:
         """
         Parse optimizer attributes for IPU dynamic to static support. Currently only support parse lr.
 
@@ -688,11 +742,11 @@ def get_lr():
 
     def set_graph_config(
         self,
-        num_ipus=1,
-        is_training=True,
-        micro_batch_size=1,
-        enable_manual_shard=False,
-    ):
+        num_ipus: int = 1,
+        is_training: bool = True,
+        micro_batch_size: int = 1,
+        enable_manual_shard: bool = False,
+    ) -> None:
         """
         Set graph configuration to the IpuStrategy instance.
 
@@ -737,11 +791,11 @@ def set_graph_config(
 
     def set_pipelining_config(
         self,
-        enable_pipelining=False,
-        batches_per_step=1,
-        enable_gradient_accumulation=False,
-        accumulation_factor=1,
-    ):
+        enable_pipelining: bool = False,
+        batches_per_step: int = 1,
+        enable_gradient_accumulation: bool = False,
+        accumulation_factor: int = 1,
+    ) -> None:
         """
         Set pipelining configuration to the IpuStrategy instance. Used to optimize the throughput performance.
 
@@ -787,7 +841,7 @@ def set_pipelining_config(
         }
         self.set_options(options)
 
-    def set_precision_config(self, enable_fp16=False):
+    def set_precision_config(self, enable_fp16: bool = False) -> None:
         """
         Set half computation configuration to the IpuStrategy instance. Used to optimize the performance.
 
@@ -816,8 +870,12 @@ def set_precision_config(self, enable_fp16=False):
         self.set_options(options)
 
     def add_custom_op(
-        self, paddle_op, popart_op=None, domain='custom.ops', version=1
-    ):
+        self,
+        paddle_op: str,
+        popart_op: str | None = None,
+        domain: str = 'custom.ops',
+        version: int = 1,
+    ) -> None:
         """
         Add a mapping to use popart custom ops running on the IPU.
 
@@ -859,7 +917,7 @@ def add_custom_op(
         if not self.has_custom_ops:
             self.has_custom_ops = True
 
-    def set_options(self, options):
+    def set_options(self, options: _IpuStrategyOptions) -> None:
         """
         Set options from dict.
 
@@ -881,15 +939,15 @@ def set_options(self, options):
 
                 >>> ipu_strategy = static.IpuStrategy()
                 >>> options = {'num_ipus':1, 'enable_fp16': True}
-                >>> ipu_strategy.set_options(options)
+                >>> ipu_strategy.set_options(options)  # type: ignore[arg-type]
         """
         self._ipu_strategy.set_options(options)
         # check whether to recompile program with updated ipu options.
         recompile_white_list = {'lr'}
         if options.keys() - recompile_white_list:
             self.need_compile = True
 
-    def get_option(self, option):
+    def get_option(self, option: str) -> Any:
         """
         Get option.
 
@@ -914,7 +972,7 @@ def get_option(self, option):
         """
         return self._ipu_strategy.get_option(option)['value']
 
-    def enable_pattern(self, pattern):
+    def enable_pattern(self, pattern: str) -> None:
         """
         Enable PopART pattern to optimize the graph.
 
@@ -939,7 +997,7 @@ def enable_pattern(self, pattern):
         """
         self._ipu_strategy.enable_pattern(pattern)
 
-    def disable_pattern(self, pattern):
+    def disable_pattern(self, pattern: str) -> None:
         """
         Disable PopART pattern.
 
@@ -965,28 +1023,28 @@ def disable_pattern(self, pattern):
         self._ipu_strategy.disable_pattern(pattern)
 
     @property
-    def num_ipus(self):
+    def num_ipus(self) -> int:
         """
         Get the number of IPU devices from IpuStrategy instance.
         """
         return self.get_option('num_ipus')
 
     @property
-    def is_training(self):
+    def is_training(self) -> bool:
         """
         Get the boolean of training or inference from IpuStrategy instance.
         """
         return self.get_option('is_training')
 
     @property
-    def enable_pipelining(self):
+    def enable_pipelining(self) -> bool:
         """
         Get the boolean of enable pipelining or not from IpuStrategy instance.
         """
         return self.get_option('enable_pipelining')
 
     @property
-    def enable_fp16(self):
+    def enable_fp16(self) -> bool:
         """
         Get the boolean of float16 mode or not from IpuStrategy instance.
         """
@@ -1036,7 +1094,12 @@ class IpuCompiledProgram:
             ...     ipu_strategy=ipu_strategy)
     """
 
-    def __init__(self, program=None, scope=None, ipu_strategy=None):
+    def __init__(
+        self,
+        program: Program | None = None,
+        scope: _Scope | None = None,
+        ipu_strategy: IpuStrategy | None = None,
+    ) -> None:
         if not core.is_compiled_with_ipu():
             raise ValueError(
                 "Can not use this function since PaddlePaddle is not compiled with IPU"
@@ -1073,7 +1136,7 @@ def __init__(self, program=None, scope=None, ipu_strategy=None):
 
         self._backend = core.IpuBackend.get_instance()
 
-    def compile(self, feed_list, fetch_list):
+    def compile(self, feed_list: list[str], fetch_list: list[str]) -> Program:
         """
         This interface is used to compile the input Program to a program
         to run the model on the ipu.
diff --git a/python/paddle/base/executor.py b/python/paddle/base/executor.py
@@ -74,7 +74,7 @@
 InferAnalysisConfig = core.AnalysisConfig
 
 
-def global_scope() -> core.Scope:
+def global_scope() -> core._Scope:
     """
     :api_attr: Static Graph
 
@@ -96,15 +96,15 @@ def global_scope() -> core.Scope:
     return g_scope
 
 
-def _switch_scope(scope: core.Scope) -> core.Scope:
+def _switch_scope(scope: core._Scope) -> core._Scope:
     global g_scope
     ex = g_scope
     g_scope = scope
     return ex
 
 
 @signature_safe_contextmanager
-def scope_guard(scope: core.Scope) -> Generator[None, None, None]:
+def scope_guard(scope: core._Scope) -> Generator[None, None, None]:
     """
 
     This function switches scope through python `with` statement.
@@ -615,7 +615,7 @@ def _fetch_var(name, scope=None, return_numpy=True):
     Args:
         name(str): name of the variable. Typically, only persistable variables
             can be found in the scope used for running the program.
-        scope(core.Scope|None): scope object. It should be the scope where
+        scope(core._Scope|None): scope object. It should be the scope where
             you pass to Executor.run() when running your program.
             If None, global_scope() will be used. Default None.
         return_numpy(bool): whether convert the tensor to numpy.ndarray.
@@ -1704,7 +1704,7 @@ def run(
         fetch_list: str | Tensor | Sequence[str | Tensor] | None = ...,
         feed_var_name: str = ...,
         fetch_var_name: str = ...,
-        scope: core.Scope | None = ...,
+        scope: core._Scope | None = ...,
         return_numpy: Literal[True] = ...,
         use_program_cache: bool = ...,
         use_prune: bool = ...,
@@ -1718,7 +1718,7 @@ def run(
         fetch_list: str | Tensor | Sequence[str | Tensor] | None = ...,
         feed_var_name: str = ...,
         fetch_var_name: str = ...,
-        scope: core.Scope | None = ...,
+        scope: core._Scope | None = ...,
         return_numpy: Literal[False] = ...,
         use_program_cache: bool = ...,
         use_prune: bool = ...,
@@ -1732,7 +1732,7 @@ def run(
         fetch_list: str | Tensor | Sequence[str | Tensor] | None = ...,
         feed_var_name: str = ...,
         fetch_var_name: str = ...,
-        scope: core.Scope | None = ...,
+        scope: core._Scope | None = ...,
         return_numpy: bool = ...,
         use_program_cache: bool = ...,
         use_prune: bool = ...,
@@ -3183,7 +3183,7 @@ def infer_from_dataset(
         self,
         program: Program | CompiledProgram | None = None,
         dataset: DatasetBase | _FleetDatasetBase | None = None,
-        scope: core.Scope | None = None,
+        scope: core._Scope | None = None,
         thread: int = 0,
         debug: bool = False,
         fetch_list: list[Tensor] | None = None,
@@ -3261,7 +3261,7 @@ def infer_from_dataset(
     def start_heter_trainer(
         self,
         program: Program | None = None,
-        scope: core.Scope | None = None,
+        scope: core._Scope | None = None,
         debug: bool = False,
         fetch_list: list[Tensor] | None = None,
         fetch_info: list[str] | None = None,
@@ -3306,7 +3306,7 @@ def train_from_dataset(
         self,
         program: Program | CompiledProgram | None = None,
         dataset: DatasetBase | _FleetDatasetBase | None = None,
-        scope: core.Scope | None = None,
+        scope: core._Scope | None = None,
         thread: int = 0,
         debug: bool = False,
         fetch_list: list[Tensor] | None = None,
diff --git a/python/paddle/incubate/distributed/utils/io/dist_save.py b/python/paddle/incubate/distributed/utils/io/dist_save.py
@@ -97,7 +97,7 @@ def save(
 
         .. code-block:: python
 
-            >>> # doctest: +SKIP('TODO: the error will be fix in the feature')
+            >>> # doctest: +SKIP('TODO: the error will be fixed in the future')
             >>> # type: ignore
             >>> import paddle
             >>> paddle.distributed.init_process_group(backend='nccl')
