[torch/elastic] Improve the implementation of RendezvousParameters and add its unit tests. (#146)

Summary:
Pull Request resolved: pytorch/elastic#146

Pull Request resolved: pytorch#54807

Improve the implementation and the unit test coverage of `RendezvousParameters`.

Test Plan: Run the existing and newly-introduced unit tests.

Reviewed By: kiukchung

Differential Revision: D27342444

fbshipit-source-id: 88de356c0a799844a739eb9105185bb8c1acf11f

Author: cbalioglu

test/distributed/elastic/rendezvous/api_test.py

@@ -3,8 +3,9 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-import unittest
-from typing import Tuple
+
+from typing import Any, Dict, SupportsInt, Tuple, cast
+from unittest import TestCase
 
 from torch.distributed.elastic.rendezvous import (
     RendezvousHandler,
@@ -40,7 +41,7 @@ def get_run_id(self) -> str:
         return ""
 
 
-class RendezvousHandlerFactoryTest(unittest.TestCase):
+class RendezvousHandlerFactoryTest(TestCase):
     def test_double_registration(self):
         factory = RendezvousHandlerFactory()
         factory.register("mock", create_mock_rdzv_handler)
@@ -67,17 +68,171 @@ def test_create_handler(self):
         self.assertTrue(isinstance(mock_rdzv_handler, MockRendezvousHandler))
 
 
-class RendezvousParametersTest(unittest.TestCase):
-    def test_get_or_default(self):
-
-        params = RendezvousParameters(
-            backend="foobar",
-            endpoint="localhost",
-            run_id="1234",
-            min_nodes=1,
-            max_nodes=1,
-            timeout1=10,
+class RendezvousParametersTest(TestCase):
+    def setUp(self) -> None:
+        self._backend = "dummy_backend"
+        self._endpoint = "dummy_endpoint"
+        self._run_id = "dummy_run_id"
+        self._min_nodes = 3
+        self._max_nodes = 6
+        self._kwargs: Dict[str, Any] = {}
+
+    def _create_params(self) -> RendezvousParameters:
+        return RendezvousParameters(
+            backend=self._backend,
+            endpoint=self._endpoint,
+            run_id=self._run_id,
+            min_nodes=self._min_nodes,
+            max_nodes=self._max_nodes,
+            **self._kwargs,
         )
 
-        self.assertEqual(10, params.get("timeout1", 20))
-        self.assertEqual(60, params.get("timeout2", 60))
+    def test_init_initializes_params(self) -> None:
+        self._kwargs["dummy_param"] = "x"
+
+        params = self._create_params()
+
+        self.assertEqual(params.backend, self._backend)
+        self.assertEqual(params.endpoint, self._endpoint)
+        self.assertEqual(params.run_id, self._run_id)
+        self.assertEqual(params.min_nodes, self._min_nodes)
+        self.assertEqual(params.max_nodes, self._max_nodes)
+
+        self.assertEqual(params.get("dummy_param"), "x")
+
+    def test_init_initializes_params_if_min_nodes_equals_to_1(self) -> None:
+        self._min_nodes = 1
+
+        params = self._create_params()
+
+        self.assertEqual(params.min_nodes, self._min_nodes)
+        self.assertEqual(params.max_nodes, self._max_nodes)
+
+    def test_init_initializes_params_if_min_nodes_equals_to_max_nodes(self) -> None:
+        self._max_nodes = 3
+
+        params = self._create_params()
+
+        self.assertEqual(params.min_nodes, self._min_nodes)
+        self.assertEqual(params.max_nodes, self._max_nodes)
+
+    def test_init_raises_error_if_backend_is_none_or_empty(self) -> None:
+        for backend in [None, ""]:
+            with self.subTest(backend=backend):
+                self._backend = backend  # type: ignore[assignment]
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    r"^The rendezvous backend name must be a non-empty string.$",
+                ):
+                    self._create_params()
+
+    def test_init_raises_error_if_min_nodes_is_less_than_1(self) -> None:
+        for min_nodes in [0, -1, -5]:
+            with self.subTest(min_nodes=min_nodes):
+                self._min_nodes = min_nodes
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    rf"^The minimum number of rendezvous nodes \({min_nodes}\) must be greater "
+                    rf"than zero.$",
+                ):
+                    self._create_params()
+
+    def test_init_raises_error_if_max_nodes_is_less_than_min_nodes(self) -> None:
+        for max_nodes in [2, 1, -2]:
+            with self.subTest(max_nodes=max_nodes):
+                self._max_nodes = max_nodes
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    rf"^The maximum number of rendezvous nodes \({max_nodes}\) must be greater "
+                    "than or equal to the minimum number of rendezvous nodes "
+                    rf"\({self._min_nodes}\).$",
+                ):
+                    self._create_params()
+
+    def test_get_returns_none_if_key_does_not_exist(self) -> None:
+        params = self._create_params()
+
+        self.assertIsNone(params.get("dummy_param"))
+
+    def test_get_returns_default_if_key_does_not_exist(self) -> None:
+        params = self._create_params()
+
+        self.assertEqual(params.get("dummy_param", default="x"), "x")
+
+    def test_get_as_bool_returns_none_if_key_does_not_exist(self) -> None:
+        params = self._create_params()
+
+        self.assertIsNone(params.get_as_bool("dummy_param"))
+
+    def test_get_as_bool_returns_default_if_key_does_not_exist(self) -> None:
+        params = self._create_params()
+
+        self.assertTrue(params.get_as_bool("dummy_param", default=True))
+
+    def test_get_as_bool_returns_true_if_value_represents_true(self) -> None:
+        for value in ["1", "True", "tRue", "T", "t", "yEs", "Y", 1, True]:
+            with self.subTest(value=value):
+                self._kwargs["dummy_param"] = value
+
+                params = self._create_params()
+
+                self.assertTrue(params.get_as_bool("dummy_param"))
+
+    def test_get_as_bool_returns_false_if_value_represents_false(self) -> None:
+        for value in ["0", "False", "faLse", "F", "f", "nO", "N", 0, False]:
+            with self.subTest(value=value):
+                self._kwargs["dummy_param"] = value
+
+                params = self._create_params()
+
+                self.assertFalse(params.get_as_bool("dummy_param"))
+
+    def test_get_as_bool_raises_error_if_value_is_invalid(self) -> None:
+        for value in ["01", "Flse", "Ture", "g", "4", "_", "truefalse", 2, -1]:
+            with self.subTest(value=value):
+                self._kwargs["dummy_param"] = value
+
+                params = self._create_params()
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    r"^The rendezvous configuration option 'dummy_param' does not represent a "
+                    r"valid boolean value.$",
+                ):
+                    params.get_as_bool("dummy_param")
+
+    def test_get_as_int_returns_none_if_key_does_not_exist(self) -> None:
+        params = self._create_params()
+
+        self.assertIsNone(params.get_as_int("dummy_param"))
+
+    def test_get_as_int_returns_default_if_key_does_not_exist(self) -> None:
+        params = self._create_params()
+
+        self.assertEqual(params.get_as_int("dummy_param", default=5), 5)
+
+    def test_get_as_int_returns_integer_if_value_represents_integer(self) -> None:
+        for value in ["0", "-10", "5", "  4", "4  ", " 4 ", 0, -4, 3]:
+            with self.subTest(value=value):
+                self._kwargs["dummy_param"] = value
+
+                params = self._create_params()
+
+                self.assertEqual(params.get_as_int("dummy_param"), int(cast(SupportsInt, value)))
+
+    def test_get_as_int_raises_error_if_value_is_invalid(self) -> None:
+        for value in ["a", "0a", "3b", "abc"]:
+            with self.subTest(value=value):
+                self._kwargs["dummy_param"] = value
+
+                params = self._create_params()
+
+                with self.assertRaisesRegex(
+                    ValueError,
+                    r"^The rendezvous configuration option 'dummy_param' does not represent a "
+                    r"valid integer value.$",
+                ):
+                    params.get_as_int("dummy_param")

torch/distributed/elastic/rendezvous/api.py

@@ -133,16 +133,22 @@ def main():
 
 
 class RendezvousParameters:
+    """Holds the parameters to construct a `RendezvousHandler`.
+
+    Args:
+        backend:
+            The name of the backend to use to handle the rendezvous.
+        endpoint:
+            The endpoint of the rendezvous, usually in form <hostname>[:<port>].
+        run_id:
+            The id of the rendezvous.
+        min_nodes:
+            The minimum number of nodes to admit to the rendezvous.
+        max_nodes:
+            The maximum number of nodes to admit to the rendezvous.
+        **kwargs:
+            Additional parameters for the specified backend.
     """
-    The data object holding parameters to construct a ``RendezvousHandler``.
-    """
-
-    # Default timeout for the rendezvous.
-    _DEFAULT_TIMEOUT: int = 600  # 10 minutes
-
-    # Additional waiting time after reaching the minimum number of nodes
-    # in case the rendezvous is elastic (min != max).
-    _DEFAULT_LAST_CALL_TIMEOUT: int = 30  # 30 seconds
 
     def __init__(
         self,
@@ -153,25 +159,17 @@ def __init__(
         max_nodes: int,
         **kwargs,
     ):
-        """
-        Args:
-            backend: The backend that is used to register the rendezvous.
-            endpoint: The endpoint of the rendezvous. Usually it is a string in the format
-                <hostname>:<port>.
-            run_id: The id of the rendezvous.
-            min_nodes: The minimum number of nodes required to complete the rendezvous.
-            max_nodes: The maximum number of nodes that are allowed to join the rendezvous.
-            **kwargs: Additional parameters for the specified backend.
-        """
-        if backend is None:
-            raise ValueError("The backend cannot be None.")
+        if not backend:
+            raise ValueError("The rendezvous backend name must be a non-empty string.")
 
         if min_nodes < 1:
-            raise ValueError("The minimum number of nodes must be greater than zero.")
+            raise ValueError(
+                f"The minimum number of rendezvous nodes ({min_nodes}) must be greater than zero."
+            )
         if max_nodes < min_nodes:
             raise ValueError(
-                "The maximum number of nodes must be greater than"
-                " or equal to the minimum number of nodes."
+                f"The maximum number of rendezvous nodes ({max_nodes}) must be greater than or "
+                f"equal to the minimum number of rendezvous nodes ({min_nodes})."
             )
 
         self.backend = backend
@@ -181,54 +179,40 @@ def __init__(
         self.max_nodes = max_nodes
         self.config = kwargs
 
-    @property
-    def timeout(self):
-        """
-        Gets the timeout for the rendezvous.
-        """
-        return self.get_as_int("timeout", self._DEFAULT_TIMEOUT)
-
-    @property
-    def last_call_timeout(self):
-        """
-        Gets additional waiting time after reaching the minimum number of nodes
-        in case the rendezvous is elastic (min != max).
-        """
-        return self.get_as_int("last_call_timeout", self._DEFAULT_LAST_CALL_TIMEOUT)
-
     def get(self, key: str, default: Any = None) -> Any:
-        """
-        Returns the value for ``key`` if ``key`` exists, else ``default``.
-        """
+        """Returns the value for `key` if `key` exists, else `default`."""
         return self.config.get(key, default)
 
     def get_as_bool(self, key: str, default: Optional[bool] = None) -> Optional[bool]:
-        """
-        Returns the value for ``key`` as a ``bool`` if ``key`` exists.
-        """
-        val = self.get(key, default)
-        if val is None:
-            return val
-        if isinstance(val, int) or isinstance(val, bool):
-            return True if val else False
-        if isinstance(val, str):
-            return val.lower() in ["1", "true", "t", "yes", "y"]
+        """Returns the value for `key` as a `bool`."""
+        value = self.get(key, default)
+        if value is None or isinstance(value, bool):
+            return value
+        if isinstance(value, int):
+            if value == 1:
+                return True
+            if value == 0:
+                return False
+        elif isinstance(value, str):
+            if value.lower() in ["1", "true", "t", "yes", "y"]:
+                return True
+            if value.lower() in ["0", "false", "f", "no", "n"]:
+                return False
         raise ValueError(
-            f"The '{key}' rendezvous config does not represent a valid boolean value."
+            f"The rendezvous configuration option '{key}' does not represent a valid boolean value."
         )
 
     def get_as_int(self, key: str, default: Optional[int] = None) -> Optional[int]:
-        """
-        Returns the value for ``key`` as an ``int`` if ``key`` exists.
-        """
-        val = self.get(key, default)
-        if val is None:
-            return val
+        """Returns the value for `key` as an `int`."""
+        value = self.get(key, default)
+        if value is None:
+            return value
         try:
-            return int(val)
+            return int(value)
         except ValueError:
             raise ValueError(
-                f"The '{key}' rendezvous config does not represent a valid integer value."
+                f"The rendezvous configuration option '{key}' does not represent a valid integer "
+                "value."
             )
 
 
@@ -269,15 +253,17 @@ def create_handler(self, params: RendezvousParameters) -> RendezvousHandler:
             creator = self._registry[params.backend]
         except KeyError:
             raise ValueError(
-                f"The rendezvous backend '{params.backend}' is not registered. Did you forget to call {self.register.__name__}?"
+                f"The rendezvous backend '{params.backend}' is not registered. Did you forget "
+                f"to call {self.register.__name__}?"
             )
 
         handler = creator(params)
 
         # Do some sanity check.
         if handler.get_backend() != params.backend:
             raise RuntimeError(
-                f"The rendezvous handler backend '{handler.get_backend()}' does not match the requested backend '{params.backend}'."
+                f"The rendezvous handler backend '{handler.get_backend()}' does not match the "
+                f"requested backend '{params.backend}'."
             )
 
         return handler

torch/distributed/elastic/rendezvous/etcd_rendezvous.py

@@ -54,6 +54,13 @@ class EtcdRendezvousRetryImmediately(Exception):
     pass
 
 
+# Default timeout for the rendezvous.
+_DEFAULT_TIMEOUT: int = 600  # 10 minutes
+
+# Additional waiting time after reaching the minimum number of nodes
+# in case the rendezvous is elastic (min != max).
+_DEFAULT_LAST_CALL_TIMEOUT: int = 30  # 30 seconds
+
 # Various constants used internally in EtcdRendezvous
 CONST_ETCD_SETUP_TTL = 5
 CONST_ETCD_FROZEN_TTL = 10
@@ -1252,7 +1259,7 @@ def create_rdzv_handler(params: RendezvousParameters) -> RendezvousHandler:
         run_id=params.run_id,
         num_min_workers=params.min_nodes,
         num_max_workers=params.max_nodes,
-        timeout=params.timeout,
-        last_call_timeout=params.last_call_timeout,
+        timeout=params.get_as_int("timeout", _DEFAULT_TIMEOUT),
+        last_call_timeout=params.get_as_int("last_call_timeout", _DEFAULT_LAST_CALL_TIMEOUT),
     )
     return EtcdRendezvousHandler(rdzv_impl=rdzv)