[Typing][C-43, C-44][BUAA] Add type annotations for python/paddle/distributed/fleet/* (#67469)

Whsjrczr · web-flow · commit a6282ef91502 · 2024-08-21T15:03:41.000+08:00
* 2 api

* update broadcast tensor

* update 2 files

* undo

* update

* update

* update sequence
diff --git a/python/paddle/distributed/fleet/base/util_factory.py b/python/paddle/distributed/fleet/base/util_factory.py
@@ -11,6 +11,8 @@
 # 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.
+from __future__ import annotations
+
 """Fleet Utils."""
 """distributed operations"""
 """basic collective operations in python"""
@@ -20,6 +22,7 @@
 import re
 import subprocess
 from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
 from google.protobuf import text_format
@@ -33,6 +36,18 @@
 from ..utils.fs import FS
 from .graphviz import GraphPreviewGenerator
 
+if TYPE_CHECKING:
+    import numpy.typing as npt
+
+    from paddle import Tensor
+    from paddle._typing import NestedNumbericSequence
+    from paddle.base.framework import Block
+    from paddle.distributed.fleet.base.distributed_strategy import (
+        DistributedStrategy,
+    )
+    from paddle.distributed.fleet.base.role_maker import PaddleCloudRoleMaker
+
+
 __all__ = []
 
 
@@ -47,23 +62,28 @@ def _create_util(self, context=None):
 
 
 class UtilBase:
-    def __init__(self):
-        self.role_maker = None
-        self.dist_strategy = None
+    def __init__(self) -> None:
+        self.role_maker: PaddleCloudRoleMaker | None = None
+        self.dist_strategy: DistributedStrategy | None = None
 
-    def _set_strategy(self, dist_strategy):
+    def _set_strategy(self, dist_strategy: DistributedStrategy | None) -> None:
         self.dist_strategy = dist_strategy
 
-    def _set_role_maker(self, role_maker):
+    def _set_role_maker(self, role_maker: PaddleCloudRoleMaker | None) -> None:
         self.role_maker = role_maker
 
-    def _set_file_system(self, fs_client):
+    def _set_file_system(self, fs_client: FS) -> None:
         assert isinstance(
             fs_client, FS
         ), "fs_client must be the instance of paddle.distributed.fleet.utils.FS"
         self.fs_client = fs_client
 
-    def all_reduce(self, input, mode="sum", comm_world="worker"):
+    def all_reduce(
+        self,
+        input: NestedNumbericSequence | npt.NDArray[Any],
+        mode: Literal["sum", "min", "max"] = "sum",
+        comm_world: Literal["worker", "server", "all"] = "worker",
+    ) -> npt.NDArray[Any] | None:
         """
         All reduce `input` between specified collection. This is a distributed API.
 
@@ -113,7 +133,9 @@ def all_reduce(self, input, mode="sum", comm_world="worker"):
             input = list(input)
         return self.role_maker._all_reduce(input, mode, comm_world)
 
-    def barrier(self, comm_world="worker"):
+    def barrier(
+        self, comm_world: Literal["worker", "server", "all"] = "worker"
+    ) -> None:
         """
         Barrier between specified collection.
 
@@ -154,7 +176,11 @@ def barrier(self, comm_world="worker"):
         """
         self.role_maker._barrier(comm_world)
 
-    def all_gather(self, input, comm_world="worker"):
+    def all_gather(
+        self,
+        input: float,
+        comm_world: Literal["worker", "server", "all"] = "worker",
+    ) -> list[float]:
         """
         All gather `input` between specified collection.
 
@@ -202,13 +228,13 @@ def all_gather(self, input, comm_world="worker"):
 
         return self.role_maker._all_gather(input, comm_world)
 
-    def _broadcast(self):
+    def _broadcast(self) -> None:
         pass
 
-    def _scatter(self):
+    def _scatter(self) -> None:
         pass
 
-    def get_heter_file_shard(self, files):
+    def get_heter_file_shard(self, files: list[str]) -> list[str]:
         if not isinstance(files, list):
             raise TypeError("files should be a list of file need to be read.")
         trainers = self.role_maker._worker_num()
@@ -228,7 +254,7 @@ def get_heter_file_shard(self, files):
 
         return trainer_files[trainer_id]
 
-    def get_file_shard(self, files):
+    def get_file_shard(self, files: list[str]) -> list[str]:
         """
         Split files before distributed training, and return filelist assigned to the current trainer.
 
@@ -287,7 +313,7 @@ def get_file_shard(self, files):
 
         return trainer_files[trainer_id]
 
-    def print_on_rank(self, message, rank_id):
+    def print_on_rank(self, message: str, rank_id: int) -> None:
         """
         Worker of rank `rank_id` print some message.
 
@@ -319,15 +345,20 @@ def print_on_rank(self, message, rank_id):
             return
         print(message)
 
-    def _save_program(self, program, model_filename='__model__', is_text=False):
+    def _save_program(
+        self,
+        program: Program,
+        model_filename: str = '__model__',
+        is_text: bool = False,
+    ) -> None:
         if is_text:
             with open(model_filename, "w") as f:
                 f.write(str(program))
         else:
             with open(model_filename, "wb") as f:
                 f.write(program.desc.serialize_to_string())
 
-    def _load_program(self, path, is_text):
+    def _load_program(self, path: str, is_text: bool) -> Program:
         def load_program_binary(path):
             """load program from binary string file"""
             with open(path, "rb") as f:
@@ -348,15 +379,19 @@ def load_program_text(path):
         else:
             return load_program_binary(path)
 
-    def _program_type_trans(self, prog_dir, prog_fn, is_text):
+    def _program_type_trans(
+        self, prog_dir: str, prog_fn: str, is_text: bool
+    ) -> str:
         prog = self._load_program(os.path.join(prog_dir, prog_fn), is_text)
         prog_out_fn = prog_fn + ".bin" if is_text else prog_fn + ".pbtxt"
         self._save_program(
             prog, os.path.join(prog_dir, prog_out_fn), 1 - is_text
         )
         return prog_out_fn
 
-    def _visualize_graphviz(self, program, output_dir, output_filename):
+    def _visualize_graphviz(
+        self, program: Program, output_dir: str, output_filename: str
+    ) -> None:
         block = program.global_block()
         dot_path = os.path.join(output_dir, output_filename + '.dot')
         pdf_path = os.path.join(output_dir, output_filename + '.pdf')
@@ -370,7 +405,7 @@ def _visualize_graphviz(self, program, output_dir, output_filename):
         )
         p.wait()
 
-    def _proto_check(self, config):
+    def _proto_check(self, config: Any) -> bool:
         train_prog = self._load_program(
             config.train_prog_path, config.is_text_train_program
         )
@@ -418,7 +453,9 @@ def _proto_check(self, config):
                 is_match = False
         return is_match
 
-    def _params_check(self, config):
+    def _params_check(
+        self, config: Any
+    ) -> list[Tensor] | list[npt.NDArray[Any]] | Literal[False]:
         def feed_gen(batch_size, feeded_vars_dims, feeded_vars_filelist):
             def reader(batch_size, fn, dim):
                 data = []
@@ -673,7 +710,9 @@ def check_not_expected_ops(prog, not_expected_op_types):
             return results
 
 
-def draw_block_graphviz(block, highlights=None, path="./temp.dot"):
+def draw_block_graphviz(
+    block: Block, highlights: list[str] | None = None, path: str = "./temp.dot"
+) -> None:
     '''
     Generate a debug graph for block.
     Args:
@@ -684,7 +723,7 @@ def draw_block_graphviz(block, highlights=None, path="./temp.dot"):
     protostr = block.desc.serialize_to_string()
     desc = framework_pb2.BlockDesc.FromString(bytes(protostr))
 
-    def need_highlight(name):
+    def need_highlight(name: str) -> bool:
         if highlights is None:
             return False
         for pattern in highlights:
diff --git a/python/paddle/distributed/fleet/data_generator/data_generator.py b/python/paddle/distributed/fleet/data_generator/data_generator.py
@@ -11,8 +11,13 @@
 # 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.
+from __future__ import annotations
 
 import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
 
 __all__ = []
 
@@ -230,7 +235,10 @@ def local_iter():
 # add more generalized DataGenerator that can adapt user-defined slot
 # for example, [(name, float_list), (name, str_list), (name, int_list)]
 class MultiSlotStringDataGenerator(DataGenerator):
-    def _gen_str(self, line):
+    def _gen_str(
+        self,
+        line: Sequence[tuple[str, list[str]]],
+    ) -> str:
         '''
         Further processing the output of the process() function rewritten by
         user, outputting data that can be directly read by the MultiSlotDataFeed,
@@ -275,7 +283,10 @@ def _gen_str(self, line):
 
 
 class MultiSlotDataGenerator(DataGenerator):
-    def _gen_str(self, line):
+    def _gen_str(
+        self,
+        line: Sequence[tuple[str, list[float]]],
+    ) -> str:
         '''
         Further processing the output of the process() function rewritten by
         user, outputting data that can be directly read by the MultiSlotDataFeed,
