[FileStore] Implemented numKeys and Added Tests (pytorch#49556)

Summary:
Pull Request resolved: pytorch#49556

Implemented the missing Store functionality (specifically numKeys) in the FileStore.

Test Plan: Added both C++ and Python tests to verify functionality.

Reviewed By: jiayisuse

Differential Revision: D25619001

fbshipit-source-id: 9146d0da9e0903622be3035880f619bbb2cc3891

test/distributed/test_c10d.py

@@ -225,6 +225,7 @@ def _test_set_get(self, fs):
         fs.add("key3", 4)
         fs.add("key3", 5)
         fs.add("key3", 6)
+        self.assertEqual(fs.num_keys(), self.num_keys_total)
         self.assertEqual(b"6", fs.get("key"))
         self.assertEqual(b"value0", fs.get("key0"))
         self.assertEqual(b"value1", fs.get("key1"))
@@ -234,6 +235,14 @@ def _test_set_get(self, fs):
     def test_set_get(self):
         self._test_set_get(self._create_store())
 
+    # This is the number of keys used in test_set_get. Adding this as a class
+    # property instead of hardcoding in the test since some Store
+    # implementations will have differing number of keys. In the base case,
+    # there will be 5 keys: key, key0, key1, key2, key3.
+    @property
+    def num_keys_total(self):
+        return 5
+
 
 class FileStoreTest(TestCase, StoreTestBase):
     def setUp(self):
@@ -296,6 +305,12 @@ def test_address_already_in_use(self):
             store1 = c10d.TCPStore(addr, port, 1, True)  # noqa: F841
             store2 = c10d.TCPStore(addr, port, 1, True)  # noqa: F841
 
+    # The TCPStore has 6 keys in test_set_get. It contains the 5 keys added by
+    # the user and one additional key used for coordinate all the workers.
+    @property
+    def num_keys_total(self):
+        return 6
+
     def _test_numkeys_delkeys(self, fs):
         # We start off with one init key in the store to coordinate workers
         self.assertEqual(fs.num_keys(), 1)
@@ -334,6 +349,13 @@ def setUp(self):
     def _create_store(self):
         return c10d.PrefixStore(self.prefix, self.tcpstore)
 
+    # The PrefixTCPStore has 6 keys in test_set_get. It contains the 5 keys
+    # added by the user and one additional key used for coordinate all the
+    # workers.
+    @property
+    def num_keys_total(self):
+        return 6
+
 
 class MyPythonStore(c10d.Store):
     def __init__(self):

torch/lib/c10d/FileStore.cpp

@@ -355,7 +355,11 @@ int64_t FileStore::add(const std::string& key, int64_t value) {
 }
 
 int64_t FileStore::getNumKeys() {
-  TORCH_CHECK(false, "getNumKeys not implemented for FileStore");
+  std::unique_lock<std::mutex> l(activeFileOpLock_);
+  File file(path_, O_RDONLY, timeout_);
+  auto lock = file.lockShared();
+  pos_ = refresh(file, pos_, cache_);
+  return cache_.size();
 }
 
 bool FileStore::deleteKey(const std::string& /* unused */) {

torch/lib/c10d/test/FileStoreTest.cpp

@@ -49,13 +49,19 @@ void testGetSet(std::string path, std::string prefix = "") {
     c10d::test::check(store, "key0", "value0");
     c10d::test::check(store, "key1", "value1");
     c10d::test::check(store, "key2", "value2");
+    auto numKeys = fileStore->getNumKeys();
+    EXPECT_EQ(numKeys, 3);
   }
 
   // Perform get on new instance
   {
     auto fileStore = c10::make_intrusive<c10d::FileStore>(path, 2);
     c10d::PrefixStore store(prefix, fileStore);
     c10d::test::check(store, "key0", "value0");
+    auto numKeys = fileStore->getNumKeys();
+    // There will be 4 keys since we still use the same underlying file as the
+    // other store above.
+    EXPECT_EQ(numKeys, 4);
   }
 }