[Typing][C-15] Add type annotations for python/paddle/distributed/communication/batch_isend_irecv.py (#66572)

enkilee · web-flow · commit 504cab3b4fa6 · 2024-07-27T21:07:36.000+08:00
diff --git a/python/paddle/distributed/communication/batch_isend_irecv.py b/python/paddle/distributed/communication/batch_isend_irecv.py
@@ -11,8 +11,10 @@
 # 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 contextlib
+from typing import TYPE_CHECKING, Callable
 
 import paddle.distributed as dist
 from paddle import framework
@@ -21,6 +23,13 @@
     _warn_cur_rank_not_in_group,
 )
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+    from paddle.base.core import task
+    from paddle.distributed import Group
+
+    _P2POpType = Callable[[Tensor, int, Group], task]
+
 
 class P2POp:
     """
@@ -61,7 +70,18 @@ class P2POp:
 
     """
 
-    def __init__(self, op, tensor, peer, group=None):
+    op: _P2POpType
+    tensor: Tensor
+    peer: int
+    group: Group | None
+
+    def __init__(
+        self,
+        op: _P2POpType,
+        tensor: Tensor,
+        peer: int,
+        group: Group | None = None,
+    ) -> None:
         if op not in [dist.isend, dist.irecv]:
             raise RuntimeError(
                 "Invalid ``op`` function. Expected ``op`` "
@@ -107,7 +127,7 @@ def _check_p2p_op_list(p2p_op_list):
         raise RuntimeError("All groups need to use the same backend.")
 
 
-def batch_isend_irecv(p2p_op_list):
+def batch_isend_irecv(p2p_op_list: list[P2POp]) -> list[task]:
     """
     Send or Receive a batch of tensors asynchronously and return a list of requests.
 
