Update on "[TCPStore] Threw value error if passing world_size=0 to TCPStore"

Andrew Gu · Andrew Gu · commit a4447f3567f1 · 2024-10-11T10:12:24.000-07:00
cc XilunWu H-Huang kwen2501 wanchaol fegin fduwjj wz337 wconstab d4l3k c-p-i-o

[ghstack-poisoned]
diff --git a/test/distributed/test_store.py b/test/distributed/test_store.py
@@ -516,6 +516,11 @@ def test_store_timeout_on_missing_clients(self):
             use_libuv=self._use_libuv,
         )
 
+    @skip_if_win32()
+    def test_world_size_0_raises(self):
+        with self.assertRaisesRegex(ValueError, "TCPStore world size cannot be 0"):
+            dist.TCPStore("localhost", 0, world_size=0, is_master=False)
+
 
 class LibUvTCPStoreTest(TCPStoreTest):
     _use_libuv = True
diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp
@@ -1534,9 +1534,9 @@ Example::
                       bool useLibUV) {
             std::optional<std::size_t> numWorkers = std::nullopt;
             if (worldSize.has_value() && worldSize.value() > -1) {
-                if (worldSize.value() == 0) {
-                    throw py::value_error("TCPStore world size cannot be 0");
-                }
+              if (worldSize.value() == 0) {
+                throw py::value_error("TCPStore world size cannot be 0");
+              }
               numWorkers = static_cast<std::size_t>(worldSize.value());
             }
 
