Add c10d filestore functionality to the current c10d_rendezvous_backend (#59719)

SirLegolot · facebook-github-bot · commit 2c5db9a40a88 · 2021-06-16T12:13:36.000-07:00
Summary: Pull Request resolved: #59719 Added filestore functionality to the c10d backend. FileStore will create a temporary file in the /tmp directory to use if it is selected as the store type. Appropriate tests were added as well. FileStore was modified to expose the path field for testing. It was also modified so that the numWorkers field in the constructor is optional (defaulting to -1). A negative value indicates there is not a fixed number of workers. In this case, the file is not attempted to be cleaned at the end. Test Plan: Unit tests for creating a c10d backend with filestore and simple error handling. Reviewed By: cbalioglu, H-Huang Differential Revision: D28997436 fbshipit-source-id: 24c9b2c9b13ea6c947e8b1207beda892bdca2217
diff --git a/test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py b/test/distributed/elastic/rendezvous/c10d_rendezvous_backend_test.py
@@ -4,14 +4,20 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+import tempfile
+
 from base64 import b64encode
 from datetime import timedelta
-from typing import ClassVar, cast
-from unittest import TestCase
+from typing import ClassVar, cast, Callable
+from unittest import TestCase, mock
 
-from torch.distributed import TCPStore
+from torch.distributed import TCPStore, FileStore
 
-from torch.distributed.elastic.rendezvous import RendezvousConnectionError, RendezvousParameters
+from torch.distributed.elastic.rendezvous import (
+    RendezvousConnectionError,
+    RendezvousParameters,
+    RendezvousError)
 from torch.distributed.elastic.rendezvous.c10d_rendezvous_backend import (
     C10dRendezvousBackend,
     create_backend,
@@ -20,7 +26,7 @@
 from rendezvous_backend_test import RendezvousBackendTestMixin
 
 
-class C10dRendezvousBackendTest(TestCase, RendezvousBackendTestMixin):
+class TCPStoreBackendTest(TestCase, RendezvousBackendTestMixin):
     _store: ClassVar[TCPStore]
 
     @classmethod
@@ -36,9 +42,31 @@ def setUp(self) -> None:
     def _corrupt_state(self) -> None:
         self._store.set("torch.rendezvous.dummy_run_id", "non_base64")
 
+class FileStoreBackendTest(TestCase, RendezvousBackendTestMixin):
+    _store: ClassVar[FileStore]
+
+    def setUp(self) -> None:
+        _, path = tempfile.mkstemp()
+        self._path = path
+
+        # Currently, filestore doesn't implement a delete_key method, so a new
+        # filestore has to be initialized for every test in order to have a
+        # clean slate.
+        self._store = FileStore(path)
+        self._backend = C10dRendezvousBackend(self._store, "dummy_run_id")
+
+    def tearDown(self) -> None:
+        os.remove(self._path)
+
+    def _corrupt_state(self) -> None:
+        self._store.set("torch.rendezvous.dummy_run_id", "non_base64")
+
 
 class CreateBackendTest(TestCase):
     def setUp(self) -> None:
+        # For testing, the default parameters used are for tcp. If a test
+        # uses parameters for file store, we set the self._params to
+        # self._params_filestore.
         self._params = RendezvousParameters(
             backend="dummy_backend",
             endpoint="localhost:29300",
@@ -50,43 +78,83 @@ def setUp(self) -> None:
             read_timeout="10",
         )
 
+        _, tmp_path = tempfile.mkstemp()
+
+        # Parameters for filestore testing.
+        self._params_filestore = RendezvousParameters(
+            backend="dummy_backend",
+            endpoint=tmp_path,
+            run_id="dummy_run_id",
+            min_nodes=1,
+            max_nodes=1,
+            store_type="fIlE",
+        )
+        self._expected_endpoint_file = tmp_path
+        self._expected_temp_dir = tempfile.gettempdir()
+
         self._expected_endpoint_host = "localhost"
         self._expected_endpoint_port = 29300
         self._expected_store_type = TCPStore
         self._expected_read_timeout = timedelta(seconds=10)
 
-    def test_create_backend_returns_backend(self) -> None:
+    def tearDown(self) -> None:
+        os.remove(self._expected_endpoint_file)
+
+
+    def _run_test_with_store(self, store_type: str, test_to_run: Callable):
+        """
+        Use this function to specify the store type to use in a test. If
+        not used, the test will default to TCPStore.
+        """
+        if store_type == "file":
+            self._params = self._params_filestore
+            self._expected_store_type = FileStore
+            self._expected_read_timeout = timedelta(seconds=300)
+
+        test_to_run()
+
+    def _assert_create_backend_returns_backend(self) -> None:
         backend, store = create_backend(self._params)
 
         self.assertEqual(backend.name, "c10d")
 
         self.assertIsInstance(store, self._expected_store_type)
 
-        tcp_store = cast(TCPStore, store)
-
-        self.assertEqual(tcp_store.host, self._expected_endpoint_host)  # type: ignore[attr-defined]
-        self.assertEqual(tcp_store.port, self._expected_endpoint_port)  # type: ignore[attr-defined]
-        self.assertEqual(tcp_store.timeout, self._expected_read_timeout)  # type: ignore[attr-defined]
+        typecast_store = cast(self._expected_store_type, store)
+        self.assertEqual(typecast_store.timeout, self._expected_read_timeout)  # type: ignore[attr-defined]
+        if (self._expected_store_type == TCPStore):
+            self.assertEqual(typecast_store.host, self._expected_endpoint_host)  # type: ignore[attr-defined]
+            self.assertEqual(typecast_store.port, self._expected_endpoint_port)  # type: ignore[attr-defined]
+        if (self._expected_store_type == FileStore):
+            if self._params.endpoint:
+                self.assertEqual(typecast_store.path, self._expected_endpoint_file)  # type: ignore[attr-defined]
+            else:
+                self.assertTrue(typecast_store.path.startswith(self._expected_temp_dir))  # type: ignore[attr-defined]
 
         backend.set_state(b"dummy_state")
 
         state = store.get("torch.rendezvous." + self._params.run_id)
 
         self.assertEqual(state, b64encode(b"dummy_state"))
 
+    def test_create_backend_returns_backend(self) -> None:
+        for store_type in ["tcp", "file"]:
+            with self.subTest(store_type=store_type):
+                self._run_test_with_store(store_type, self._assert_create_backend_returns_backend)
+
     def test_create_backend_returns_backend_if_is_host_is_false(self) -> None:
         store = TCPStore(  # type: ignore[call-arg] # noqa: F841
             self._expected_endpoint_host, self._expected_endpoint_port, is_master=True
         )
 
         self._params.config["is_host"] = "false"
 
-        self.test_create_backend_returns_backend()
+        self._assert_create_backend_returns_backend()
 
     def test_create_backend_returns_backend_if_is_host_is_not_specified(self) -> None:
         del self._params.config["is_host"]
 
-        self.test_create_backend_returns_backend()
+        self._assert_create_backend_returns_backend()
 
     def test_create_backend_returns_backend_if_is_host_is_not_specified_and_store_already_exists(
         self,
@@ -97,26 +165,35 @@ def test_create_backend_returns_backend_if_is_host_is_not_specified_and_store_al
 
         del self._params.config["is_host"]
 
-        self.test_create_backend_returns_backend()
+        self._assert_create_backend_returns_backend()
 
     def test_create_backend_returns_backend_if_endpoint_port_is_not_specified(self) -> None:
         self._params.endpoint = self._expected_endpoint_host
 
         self._expected_endpoint_port = 29400
 
-        self.test_create_backend_returns_backend()
+        self._assert_create_backend_returns_backend()
+
+    def test_create_backend_returns_backend_if_endpoint_file_is_not_specified(self) -> None:
+        self._params_filestore.endpoint = ""
+
+        self._run_test_with_store("file", self._assert_create_backend_returns_backend)
 
     def test_create_backend_returns_backend_if_store_type_is_not_specified(self) -> None:
         del self._params.config["store_type"]
 
-        self.test_create_backend_returns_backend()
+        self._expected_store_type = TCPStore
+        if (not self._params.get("read_timeout")):
+            self._expected_read_timeout = timedelta(seconds=60)
+
+        self._assert_create_backend_returns_backend()
 
     def test_create_backend_returns_backend_if_read_timeout_is_not_specified(self) -> None:
         del self._params.config["read_timeout"]
 
         self._expected_read_timeout = timedelta(seconds=60)
 
-        self.test_create_backend_returns_backend()
+        self._assert_create_backend_returns_backend()
 
     def test_create_backend_raises_error_if_store_is_unreachable(self) -> None:
         self._params.config["is_host"] = "false"
@@ -146,7 +223,7 @@ def test_create_backend_raises_error_if_store_type_is_invalid(self) -> None:
         self._params.config["store_type"] = "dummy_store_type"
 
         with self.assertRaisesRegex(
-            ValueError, r"^The store type must be 'tcp'. Other store types are not supported yet.$"
+            ValueError, r"^Invalid store type given. Currently only supports file and tcp.$"
         ):
             create_backend(self._params)
 
@@ -159,3 +236,25 @@ def test_create_backend_raises_error_if_read_timeout_is_invalid(self) -> None:
                     ValueError, r"^The read timeout must be a positive integer.$"
                 ):
                     create_backend(self._params)
+
+    @mock.patch("tempfile.mkstemp")
+    def test_create_backend_raises_error_if_tempfile_creation_fails(self, tempfile_mock) -> None:
+        tempfile_mock.side_effect = OSError("test error")
+        # Set the endpoint to empty so it defaults to creating a temp file
+        self._params_filestore.endpoint = ""
+        with self.assertRaisesRegex(
+            RendezvousError,
+            r"The file creation for C10d store has failed. See inner exception for details."
+        ):
+            create_backend(self._params_filestore)
+
+    @mock.patch("torch.distributed.elastic.rendezvous.c10d_rendezvous_backend.FileStore")
+    def test_create_backend_raises_error_if_file_path_is_invalid(self, filestore_mock) -> None:
+        filestore_mock.side_effect = RuntimeError("test error")
+        self._params_filestore.endpoint = "bad file path"
+        with self.assertRaisesRegex(
+            RendezvousConnectionError,
+            r"^The connection to the C10d store has failed. See inner exception for "
+            r"details.$",
+        ):
+            create_backend(self._params_filestore)
diff --git a/torch/_C/_distributed_c10d.pyi b/torch/_C/_distributed_c10d.pyi
@@ -128,7 +128,7 @@ class Store:
     def wait(self, keys: List[str], timeout: timedelta): ...
 
 class FileStore(Store):
-    def __init__(self, path: str, numWorkers: int): ...
+    def __init__(self, path: str, numWorkers: int = ...): ...
 
 class HashStore(Store):
     def __init__(self): ...
diff --git a/torch/csrc/distributed/c10d/init.cpp b/torch/csrc/distributed/c10d/init.cpp
@@ -807,7 +807,7 @@ A store implementation that uses a file to store the underlying key-value pairs.
 
 Arguments:
     file_name (str): path of the file in which to store the key-value pairs
-    world_size (int): The total number of processes using the store
+    world_size (int, optional): The total number of processes using the store. Default is -1 (a negative value indicates a non-fixed number of store users).
 
 Example::
     >>> import torch.distributed as dist
@@ -818,7 +818,13 @@ Example::
     >>> store2.get("first_key")
 
       )")
-      .def(py::init<const std::string&, int>());
+      .def(py::init<const std::string&, int>(),
+           py::arg("file_name"),
+           py::arg("world_size") = -1)
+      .def_property_readonly(
+          "path",
+          &::c10d::FileStore::getPath,
+          R"(Gets the path of the file used by FileStore to store key-value pairs.)");
 
 #ifndef _WIN32
   intrusive_ptr_class_<::c10d::HashStore>(
@@ -854,7 +860,7 @@ the server to establish a connection.
 Arguments:
     host_name (str): The hostname or IP Address the server store should run on.
     port (int): The port on which the server store should listen for incoming requests.
-    world_size (int, optional): The total number of store users (number of clients + 1 for the server). Default is -1 (a negative value indicates an non-fixed number of store users).
+    world_size (int, optional): The total number of store users (number of clients + 1 for the server). Default is -1 (a negative value indicates a non-fixed number of store users).
     is_master (bool, optional): True when initializing the server store and False for client stores. Default is False.
     timeout (timedelta, optional): Timeout used by the store during initialization and for methods such as :meth:`~torch.distributed.store.get` and :meth:`~torch.distributed.store.wait`. Default is timedelta(seconds=300)
     wait_for_worker (bool, optional): Whether to wait for all the workers to connect with the server store. This is only applicable when world_size is a fixed value. Default is True.
diff --git a/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py b/torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py
@@ -7,13 +7,14 @@
 import binascii
 import logging
 import os
+import tempfile
 from base64 import b64decode, b64encode
 from datetime import timedelta
 from typing import Any, Optional, Tuple, cast
 
-from torch.distributed import Store, TCPStore
+from torch.distributed import Store, TCPStore, FileStore
 
-from .api import RendezvousConnectionError, RendezvousParameters, RendezvousStateError
+from .api import RendezvousConnectionError, RendezvousParameters, RendezvousStateError, RendezvousError
 from .dynamic_rendezvous import RendezvousBackend, Token
 from .utils import _matches_machine_hostname, parse_rendezvous_endpoint
 
@@ -166,6 +167,28 @@ def _create_tcp_store(params: RendezvousParameters) -> TCPStore:
 
     return store
 
+def _create_file_store(params: RendezvousParameters) -> FileStore:
+    # If a user specifies an endpoint, we treat it as a path to a file.
+    if params.endpoint:
+        path = params.endpoint
+    else:
+        try:
+            # The temporary file is readable and writable only by the user of
+            # this process.
+            _, path = tempfile.mkstemp()
+        except OSError as exc:
+            raise RendezvousError(
+                "The file creation for C10d store has failed. See inner exception for details."
+            ) from exc
+
+    try:
+        store = FileStore(path)
+    except (ValueError, RuntimeError) as exc:
+        raise RendezvousConnectionError(
+            "The connection to the C10d store has failed. See inner exception for details."
+        ) from exc
+
+    return store
 
 def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend, Store]:
     """Creates a new :py:class:`C10dRendezvousBackend` from the specified
@@ -174,12 +197,19 @@ def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend,
     +--------------+-----------------------------------------------------------+
     | Parameter    | Description                                               |
     +==============+===========================================================+
-    | store_type   | The type of the C10d store. As of today the only          |
-    |              | supported type is "tcp" which corresponds to              |
-    |              | :py:class:`torch.distributed.TCPStore`. Defaults to "tcp".|
+    | store_type   | The type of the C10d store. The currently supported types |
+    |              | are "tcp" and "file" which correspond to                  |
+    |              | :py:class:`torch.distributed.TCPStore` and                |
+    |              | :py:class:`torch.distributed.FileStore`, respectively.    |
+    |              | Defaults to "tcp".                                        |
     +--------------+-----------------------------------------------------------+
     | read_timeout | The read timeout, in seconds, for store operations.       |
     |              | Defaults to 60 seconds.                                   |
+    |              |                                                           |
+    |              | Note this only applies to                                 |
+    |              | :py:class:`torch.distributed.TCPStore`. It is not relevant|
+    |              | to :py:class:`torch.distributed.FileStore` which does not |
+    |              | take in timeout as a parameter.                           |
     +--------------+-----------------------------------------------------------+
     | is_host      | A boolean value indicating whether this backend instance  |
     |              | will host the C10d store. If not specified it will be     |
@@ -195,12 +225,15 @@ def create_backend(params: RendezvousParameters) -> Tuple[C10dRendezvousBackend,
     |              | the hostname or does not match the FQDN of the machine).  |
     +--------------+-----------------------------------------------------------+
     """
-    # As of today we only support TCPStore. Other store types do not have the
-    # required functionality (e.g. compare_set) yet.
+    # As of today we only support TCPStore and FileStore. Other store types do
+    # not have the required functionality (e.g. compare_set) yet.
     store_type = params.get("store_type", "tcp").strip().lower()
-    if store_type != "tcp":
-        raise ValueError("The store type must be 'tcp'. Other store types are not supported yet.")
-
-    store = _create_tcp_store(params)
+    store: Store
+    if store_type == "file":
+        store = _create_file_store(params)
+    elif store_type == "tcp":
+        store = _create_tcp_store(params)
+    else:
+        raise ValueError("Invalid store type given. Currently only supports file and tcp.")
 
     return C10dRendezvousBackend(store, params.run_id), store
diff --git a/torch/lib/c10d/FileStore.cpp b/torch/lib/c10d/FileStore.cpp
@@ -272,18 +272,16 @@ FileStore::FileStore(const std::string& path, int numWorkers)
       numWorkers_(numWorkers),
       cleanupKey_("cleanup/"),
       regularPrefix_("/") {
-  if (numWorkers_ < 1) {
-    TORCH_CHECK(false,
-        "Number of workers for FileStore should be greater than zero");
-  }
 }
 
 FileStore::~FileStore() {
   // cleanup key will be different from all rest keys since all rest keys will
   // have a regular prefix.
   auto numFinishedWorker = addHelper(cleanupKey_, 1);
-  // The last worker cleans up the file
-  if (numFinishedWorker == numWorkers_) {
+  // The last worker cleans up the file. If numWorkers was not initialized to
+  // a specific postive value (i.e. meaning that there was not a fixed number
+  // of workers), we don't attempt to clean.
+  if (numWorkers_ >= 0 && numFinishedWorker == numWorkers_) {
     // Best effort removal without checking the return
     std::remove(path_.c_str());
   }
diff --git a/torch/lib/c10d/FileStore.hpp b/torch/lib/c10d/FileStore.hpp
