dist: add list_keys to Store API (pytorch#167883)

This adds a `list` Store API and implements it for all backends.

This is intended to be used for debugging and will allow inspecting all keys in a store locally as well as remotely in the case of TCPStore.

Test plan:

```
pytest test/distributed/test_store.py
```

Pull Request resolved: pytorch#167883
Approved by: https://github.com/fduwjj

Author: d4l3k

test/distributed/test_store.py

@@ -253,6 +253,14 @@ def test_clone(self):
         a.set("foo", "bar")
         self.assertEqual(b.get("foo"), b"bar")
 
+    def test_list_keys(self):
+        a = self._create_store()
+        a.set("foo", "bar")
+        a.set("baz", "qux")
+        keys = a.list_keys()
+        self.assertIn("foo", keys)
+        self.assertIn("baz", keys)
+
     # 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,

torch/_C/_distributed_c10d.pyi

@@ -215,6 +215,7 @@ class Store:
     def queue_pop(self, key: str, block: bool = True) -> bytes: ...
     def queue_push(self, key: str, value: Union[bytes, str]) -> None: ...
     def queue_len(self, key: str) -> int: ...
+    def list_keys(self) -> list[str]: ...
 
 class FileStore(Store):
     def __init__(self, path: str, numWorkers: int = ...) -> None: ...

torch/csrc/distributed/c10d/FileStore.cpp

@@ -492,4 +492,17 @@ void FileStore::wait(
   }
 }
 
+std::vector<std::string> FileStore::listKeys() {
+  std::unique_lock<std::mutex> l(activeFileOpLock_);
+  File file(path_, O_RDONLY, timeout_);
+  auto lock = file.lockShared();
+  pos_ = refresh(file, pos_, cache_, deletePrefix_);
+  std::vector<std::string> keys;
+  keys.reserve(cache_.size());
+  for (const auto& kv : cache_) {
+    keys.push_back(kv.first.substr(regularPrefix_.size()));
+  }
+  return keys;
+}
+
 } // namespace c10d

torch/csrc/distributed/c10d/FileStore.hpp

@@ -45,6 +45,8 @@ class TORCH_API FileStore : public Store {
     return path_;
   }
 
+  std::vector<std::string> listKeys() override;
+
  protected:
   int64_t addHelper(const std::string& key, int64_t i);
 

torch/csrc/distributed/c10d/HashStore.cpp

@@ -217,4 +217,14 @@ int64_t HashStore::queueLen(const std::string& key) {
   return static_cast<int64_t>(it->second.size());
 }
 
+std::vector<std::string> HashStore::listKeys() {
+  std::unique_lock<std::mutex> lock(m_);
+  std::vector<std::string> keys;
+  keys.reserve(map_.size());
+  for (const auto& kv : map_) {
+    keys.push_back(kv.first);
+  }
+  return keys;
+}
+
 } // namespace c10d

torch/csrc/distributed/c10d/HashStore.hpp

@@ -59,6 +59,8 @@ class TORCH_API HashStore : public Store {
 
   int64_t queueLen(const std::string& key) override;
 
+  std::vector<std::string> listKeys() override;
+
  protected:
   bool checkLocked(
       const std::unique_lock<std::mutex>& lock,

torch/csrc/distributed/c10d/PrefixStore.cpp

@@ -146,4 +146,18 @@ c10::intrusive_ptr<Store> PrefixStore::getUnderlyingNonPrefixStore() {
   return store;
 }
 
+std::vector<std::string> PrefixStore::listKeys() {
+  auto keys = store_->listKeys();
+  std::vector<std::string> filteredKeys;
+  filteredKeys.reserve(keys.size());
+
+  for (auto& key : keys) {
+    if (key.find(prefix_) == 0) {
+      key = key.substr(prefix_.size() + 1);
+      filteredKeys.push_back(std::move(key));
+    }
+  }
+  return filteredKeys;
+}
+
 } // namespace c10d

torch/csrc/distributed/c10d/PrefixStore.hpp

@@ -64,6 +64,8 @@ class TORCH_API PrefixStore : public Store {
   // Recursively to fetch the store before layers of wrapping with PrefixStore.
   c10::intrusive_ptr<Store> getUnderlyingNonPrefixStore();
 
+  std::vector<std::string> listKeys() override;
+
  protected:
   std::string prefix_;
   c10::intrusive_ptr<Store> store_;

torch/csrc/distributed/c10d/Store.hpp

@@ -114,6 +114,11 @@ class TORCH_API Store : public torch::CustomClassHolder {
     C10_THROW_ERROR(NotImplementedError, "queue support is not implemented.");
   }
 
+  virtual std::vector<std::string> listKeys() {
+    C10_THROW_ERROR(
+        NotImplementedError, "listKeys support is not implemented.");
+  }
+
  protected:
   std::chrono::milliseconds timeout_;
 };

torch/csrc/distributed/c10d/TCPStore.cpp

@@ -723,6 +723,30 @@ int64_t TCPStore::queueLen(const std::string& key) {
   return client_->receiveValue<int64_t>();
 }
 
+std::vector<std::string> TCPStore::listKeys() {
+  STATIC_SCOPED_WAIT_COUNTER(pytorch.wait_counter.TCPStore__list);
+
+  const std::lock_guard<std::mutex> lock(activeOpLock_);
+
+  detail::SendBuffer buffer(*client_, detail::QueryType::LIST_KEYS);
+  buffer.flush();
+
+  auto numKeys = client_->receiveValue<int64_t>();
+  std::vector<std::string> keys;
+  keys.reserve(numKeys);
+  for (auto i = 0; i < numKeys; ++i) {
+    auto bits = client_->receiveBits();
+    std::string str(bits.begin(), bits.end());
+    if (str.find(keyPrefix_) == 0) {
+      str = str.substr(keyPrefix_.size());
+    } else {
+      continue;
+    }
+    keys.emplace_back(str);
+  }
+  return keys;
+}
+
 bool TCPStore::hasExtendedApi() const {
   return true;
 }