[Typing][C-50] Add type annotations for python/paddle/distributed/fleet/utils/ps_util.py

python/paddle/distributed/fleet/utils/ps_util.py

@@ -13,11 +13,20 @@
 # limitations under the License.
 """Parameter Server utils"""
 
+from __future__ import annotations
+
 import os
 import warnings
+from typing import TYPE_CHECKING
 
 import paddle
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base import Executor
+    from paddle.distributed.fleet.base.role_maker import RoleMakerBase
+    from paddle.framework import Program
+
 __all__ = []
 
 
@@ -26,7 +35,11 @@ class DistributedInfer:
     Utility class for distributed infer of PaddlePaddle.
     """
 
-    def __init__(self, main_program=None, startup_program=None):
+    def __init__(
+        self,
+        main_program: Program | None = None,
+        startup_program: Program | None = None,
+    ) -> None:
         if main_program:
             self.origin_main_program = main_program.clone()
         else:
@@ -43,8 +56,12 @@ def __init__(self, main_program=None, startup_program=None):
         self.sparse_table_maps = None
 
     def init_distributed_infer_env(
-        self, exe, loss, role_maker=None, dirname=None
-    ):
+        self,
+        exe: Executor,
+        loss: Tensor,
+        role_maker: RoleMakerBase | None = None,
+        dirname: str | None = None,
+    ) -> None:
         from paddle.distributed import fleet
 
         if fleet.fleet._runtime_handle is None:
@@ -112,7 +129,7 @@ def _init_dense_params(self, exe=None, dirname=None):
                 vars=need_load_vars,
             )
 
-    def get_dist_infer_program(self):
+    def get_dist_infer_program(self) -> Program:
         varname2tables = self._get_sparse_table_map()
         convert_program = self._convert_program(
             self.origin_main_program, varname2tables