[Typing][C-27,C-28,C-29,C-30,C-31,C-32,C-33] Add type annotations for python/paddle/distributed/communication/stream/*.py (#66993)

megemini · SigureMo · web-flow · commit 907427ba909f · 2024-08-05T09:59:57.000+08:00
---------

Co-authored-by: SigureMo &lt;sigure.qaq@gmail.com&gt;
diff --git a/python/paddle/distributed/communication/stream/broadcast.py b/python/paddle/distributed/communication/stream/broadcast.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from paddle import _C_ops, framework
 from paddle.base import data_feeder
 from paddle.distributed.communication.group import (
@@ -22,6 +26,11 @@
 from paddle.distributed.communication.reduce import _to_inplace_op
 from paddle.framework import in_pir_mode
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
 
 def _broadcast_in_dygraph(
     tensor, src_rank_in_group, group, sync_op, use_calc_stream
@@ -78,15 +87,21 @@ def _broadcast_in_static_mode(
     )
 
 
-def broadcast(tensor, src, group=None, sync_op=True, use_calc_stream=False):
+def broadcast(
+    tensor: Tensor,
+    src: int,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Broadcast a tensor to all devices.
 
     Args:
         tensor (Tensor): The tensor to broadcast. Support float16, float32, float64, int32, int64, int8, uint8 or bool as its data type.
         src (int, optional): Rank of the source device.
-        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
         sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
         use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
             option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
@@ -111,7 +126,7 @@ def broadcast(tensor, src, group=None, sync_op=True, use_calc_stream=False):
             >>> else:
             ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
             >>> task = dist.stream.broadcast(data, src=1, sync_op=False)
-            >>> task.wait()
+            >>> task.wait()  # type: ignore[union-attr]
             >>> out = data.numpy()
             >>> print(out)
             >>> # [[1, 2, 3], [1, 2, 3]] (2 GPUs)
diff --git a/python/paddle/distributed/communication/stream/gather.py b/python/paddle/distributed/communication/stream/gather.py
@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import warnings
+from typing import TYPE_CHECKING
 
 import paddle
 import paddle.distributed as dist
@@ -23,6 +26,13 @@
     _warn_cur_rank_not_in_group,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
 
 def _gather_in_dygraph(
     tensor, gather_list, dst_rank_in_group, group, sync_op, use_calc_stream
@@ -49,24 +59,24 @@ def _gather_in_dygraph(
 
 
 def gather(
-    tensor,
-    gather_list=None,
-    dst=0,
-    group=None,
-    sync_op=True,
-    use_calc_stream=False,
-):
+    tensor: Tensor,
+    gather_list: Sequence[Tensor] | None = None,
+    dst: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Gather tensors from all participators.
 
     Args:
         tensor (Tensor): The input Tensor. Its data type
             should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16.
-        gather_list (list): A list of Tensors to hold the gathered tensors. Every element in the list must be a Tensor whose data type
+        gather_list (list|None): A list of Tensors to hold the gathered tensors. Every element in the list must be a Tensor whose data type
             should be float16, float32, float64, int32, int64, int8, uint8, bool or bfloat16. Default value is None.
         dst (int): The dst rank id. Default value is 0.
-        group (Group, optional): The group instance return by new_group or None for global default group.
+        group (Group|None, optional): The group instance return by new_group or None for global default group.
         sync_op (bool, optional): Whether this op is a sync op. The default value is True.
         use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
             option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
diff --git a/python/paddle/distributed/communication/stream/recv.py b/python/paddle/distributed/communication/stream/recv.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from paddle import framework
 from paddle.base import data_feeder
 from paddle.distributed.communication.group import (
@@ -20,6 +24,11 @@
     _warn_cur_rank_not_in_group,
 )
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
 
 def _recv_in_dygraph(
     tensor, src_rank_in_group, group, sync_op, use_calc_stream
@@ -61,15 +70,21 @@ def _recv_in_static_mode(
     )
 
 
-def recv(tensor, src=0, group=None, sync_op=True, use_calc_stream=False):
+def recv(
+    tensor: Tensor,
+    src: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Receive a tensor from the source device.
 
     Args:
         tensor (Tensor): The tensor to receive. Support float16, float32, float64, int32, int64, int8, uint8 or bool as its data type.
         src (int, optional): Rank of the source device. If none is given, use `0` as default.
-        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
         sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
         use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
             option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
@@ -92,7 +107,7 @@ def recv(tensor, src=0, group=None, sync_op=True, use_calc_stream=False):
             >>> else:
             ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
             ...     task = dist.stream.recv(data, src=0, sync_op=False)
-            >>> task.wait()
+            >>> task.wait()  # type: ignore[union-attr]
             >>> out = data.numpy()
             >>> print(out)
             >>> # [[4, 5, 6], [4, 5, 6]] (2 GPUs)
diff --git a/python/paddle/distributed/communication/stream/reduce.py b/python/paddle/distributed/communication/stream/reduce.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from paddle import framework
 from paddle.base import data_feeder
 from paddle.distributed.communication.group import (
@@ -21,6 +25,12 @@
 )
 from paddle.distributed.communication.reduce import ReduceOp, _get_reduce_op
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+    from paddle.distributed.communication.reduce import _ReduceOp
+
 
 def _reduce_in_dygraph(
     tensor, dst_rank_in_group, op, group, sync_op, use_calc_stream
@@ -76,13 +86,13 @@ def _reduce_in_static_mode(
 
 
 def reduce(
-    tensor,
-    dst=0,
-    op=ReduceOp.SUM,
-    group=None,
-    sync_op=True,
-    use_calc_stream=False,
-):
+    tensor: Tensor,
+    dst: int = 0,
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Perform specific reduction (for example, sum, max) on a tensor across devices and send to the destination device.
@@ -92,7 +102,7 @@ def reduce(
             float16, float32, float64, int32, int64, int8, uint8 or bool as the input data type.
         dst (int, optional): Rank of the destination device. If none is given, use `0` as default.
         op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
-        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
         sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
         use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
             option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
@@ -117,7 +127,7 @@ def reduce(
             >>> else:
             ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
             >>> task = dist.stream.reduce(data, dst=0, sync_op=False)
-            >>> task.wait()
+            >>> task.wait()  # type: ignore[union-attr]
             >>> out = data.numpy()
             >>> print(out)
             >>> # [[5, 7, 9], [5, 7, 9]] (2 GPUs, out for rank 0)
diff --git a/python/paddle/distributed/communication/stream/reduce_scatter.py b/python/paddle/distributed/communication/stream/reduce_scatter.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 import paddle
 import paddle.distributed as dist
 from paddle import framework
@@ -22,6 +26,14 @@
 )
 from paddle.distributed.communication.reduce import ReduceOp, _get_reduce_op
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+    from paddle.distributed.communication.reduce import _ReduceOp
+
 
 def _reduce_scatter_tensor_in_dygraph(
     out_tensor,
@@ -102,13 +114,13 @@ def _reduce_scatter_in_static_mode(tensor, tensor_or_tensor_list, group):
 
 
 def reduce_scatter(
-    tensor,
-    tensor_or_tensor_list,
-    op=ReduceOp.SUM,
-    group=None,
-    sync_op=True,
-    use_calc_stream=False,
-):
+    tensor: Tensor,
+    tensor_or_tensor_list: Tensor | Sequence[Tensor],
+    op: _ReduceOp = ReduceOp.SUM,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Reduce, then scatter a tensor (or a tensor list) across devices.
@@ -119,7 +131,7 @@ def reduce_scatter(
         tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The input to scatter.
             If it is a tensor, it should be correctly-sized. If it is a list, it should contain correctly-sized tensors.
         op (ReduceOp.SUM|ReduceOp.MAX|ReduceOp.MIN|ReduceOp.PROD, optional): The reduction used. If none is given, use ReduceOp.SUM as default.
-        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
         sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
         use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
             option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
diff --git a/python/paddle/distributed/communication/stream/scatter.py b/python/paddle/distributed/communication/stream/scatter.py
@@ -12,7 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import warnings
+from typing import TYPE_CHECKING
 
 import paddle
 import paddle.distributed as dist
@@ -24,6 +27,13 @@
     _warn_cur_rank_not_in_group,
 )
 
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
 
 def _scatter_tensor_in_dygraph(
     out_tensor, in_tensor, src_rank_in_group, group, sync_op, use_calc_stream
@@ -131,13 +141,13 @@ def _scatter_in_static_mode(
 
 
 def scatter(
-    tensor,
-    tensor_or_tensor_list=None,
-    src=0,
-    group=None,
-    sync_op=True,
-    use_calc_stream=False,
-):
+    tensor: Tensor,
+    tensor_or_tensor_list: Tensor | Sequence[Tensor] | None = None,
+    src: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Scatter a tensor (or a tensor list) across devices.
@@ -148,7 +158,7 @@ def scatter(
         tensor_or_tensor_list (Union[Tensor, List[Tensor]]): The input to scatter (default is `None`, must be specified on the source rank).
             If it is a tensor, it should be correctly-sized. If it is a list, it should contain correctly-sized tensors.
         src (int, optional): Rank of the source device. If none is given, use `0` as default.
-        group (Group, optional): Communicate in which group. If none is given, use the global group as default.
+        group (Group|None, optional): Communicate in which group. If none is given, use the global group as default.
         sync_op (bool, optional): Indicate whether the communication is sync or not. If none is given, use true as default.
         use_calc_stream (bool, optional): Indicate whether the communication is done on calculation stream. If none is given, use false as default. This
             option is designed for high performance demand, be careful to turn it on except you are clearly know its meaning.
diff --git a/python/paddle/distributed/communication/stream/send.py b/python/paddle/distributed/communication/stream/send.py
@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
 from paddle import framework
 from paddle.base import data_feeder
 from paddle.distributed.communication.group import (
@@ -20,6 +24,11 @@
     _warn_cur_rank_not_in_group,
 )
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed.communication.group import Group
+
 
 def _send_in_dygraph(
     tensor, dst_rank_in_group, group, sync_op, use_calc_stream
@@ -60,7 +69,13 @@ def _send_in_static_mode(
     )
 
 
-def send(tensor, dst=0, group=None, sync_op=True, use_calc_stream=False):
+def send(
+    tensor: Tensor,
+    dst: int = 0,
+    group: Group | None = None,
+    sync_op: bool = True,
+    use_calc_stream: bool = False,
+) -> task | None:
     """
 
     Send a tensor to the destination device.
@@ -91,7 +106,7 @@ def send(tensor, dst=0, group=None, sync_op=True, use_calc_stream=False):
             >>> else:
             ...     data = paddle.to_tensor([[1, 2, 3], [1, 2, 3]])
             ...     task = dist.stream.recv(data, src=0, sync_op=False)
-            >>> task.wait()
+            >>> task.wait()  # type: ignore[union-attr]
             >>> out = data.numpy()
             >>> print(out)
             [[4, 5, 6], [4, 5, 6]]
