[Distributed] deleteKey support for HashStore (#46049)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46049
Adding support for the deleteKey API in the c10d HashStore.
ghstack-source-id: 113874207
Test Plan:
Added C++ tests to check whether deleteKey function works, and
whether it returns an exception for attempting to delete non-existing keys.
Reviewed By: jiayisuse
Differential Revision: D24067657
fbshipit-source-id: 4c58dab407c6ffe209585ca91aa430850261b29e
diff --git a/torch/lib/c10d/HashStore.cpp b/torch/lib/c10d/HashStore.cpp
index b08f49e..0cd8918 100644
--- a/torch/lib/c10d/HashStore.cpp
+++ b/torch/lib/c10d/HashStore.cpp
@@ -84,8 +84,10 @@
return map_.size();
}
-bool HashStore::deleteKey(const std::string& /* unused */) {
- TORCH_CHECK(false, "deleteKey not implemented for HashStore");
+bool HashStore::deleteKey(const std::string& key) {
+ std::unique_lock<std::mutex> lock(m_);
+ auto numDeleted = map_.erase(key);
+ return (numDeleted == 1);
}
bool HashStore::check(const std::vector<std::string>& keys) {
diff --git a/torch/lib/c10d/test/HashStoreTest.cpp b/torch/lib/c10d/test/HashStoreTest.cpp
index ac60bba..a16f832 100644
--- a/torch/lib/c10d/test/HashStoreTest.cpp
+++ b/torch/lib/c10d/test/HashStoreTest.cpp
@@ -21,6 +21,13 @@
c10d::test::check(store, "key2", "value2");
auto numKeys = store.getNumKeys();
EXPECT_EQ(numKeys, 3);
+ auto delSuccess = store.deleteKey("key0");
+ EXPECT_TRUE(delSuccess);
+ numKeys = store.getNumKeys();
+ EXPECT_EQ(numKeys, 2);
+ auto delFailure = store.deleteKey("badKeyName");
+ EXPECT_FALSE(delFailure);
+ EXPECT_THROW(store.get("key0"), std::runtime_error);
}
// get() waits up to timeout_.
