Fix typing errors in torch.distributed.distributed_c10d.* (#47532)

Summary: Pull Request resolved: #47532

Test Plan: Imported from OSS

Reviewed By: walterddr

Differential Revision: D24952501

Pulled By: xuzhao9

fbshipit-source-id: 9b2dd1069eb1729c24be00f46da60d6a0439a8da

mypy.ini

@@ -59,6 +59,9 @@ ignore_errors = True
 [mypy-torch.distributed.rpc.*]
 ignore_errors = False
 
+[mypy-torch.distributed.distributed_c10d.*]
+ignore_errors = False
+
 [mypy-torch.testing._internal.hypothesis_utils.*]
 ignore_errors = True
 

torch/_C/_distributed_c10d.pyi

@@ -5,6 +5,9 @@ from datetime import timedelta
 
 # This module is defined in torch/csrc/distributed/c10d/init.cpp
 
+_DEFAULT_FIRST_BUCKET_BYTES: int
+_DEFAULT_NO_TIMEOUT: timedelta
+
 class BuiltinCommHookType(Enum):
     ALLREDUCE = ...
     FP16_COMPRESS = ...
@@ -122,7 +125,7 @@ class Work:
     def is_completed(self) -> bool: ...
     def is_success(self) -> bool: ...
     def exception(self) -> Any: ...
-    def wait(self, timeout: timedelta) -> bool: ...
+    def wait(self, timeout: timedelta = _DEFAULT_NO_TIMEOUT) -> bool: ...
     def source_rank(self) -> int: ...
     def _source_rank(self) -> int: ...
     def result(self) -> List[Tensor]: ...
@@ -348,5 +351,3 @@ def _broadcast_coalesced(
     src: int,
 ): ...
 def _test_python_store(store: Store): ...
-
-_DEFAULT_FIRST_BUCKET_BYTES: int

torch/csrc/distributed/c10d/init.cpp

@@ -1217,6 +1217,7 @@ that adds a prefix to each key inserted to the store.
       py::call_guard<py::gil_scoped_release>());
 
   module.attr("_DEFAULT_FIRST_BUCKET_BYTES") = ::c10d::kDefaultFirstBucketBytes;
+  module.attr("_DEFAULT_NO_TIMEOUT") = py::cast(kNoTimeout);
 
   Py_RETURN_TRUE;
 }