OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2026-01-15 12:15:51 +00:00
Code Issues Actions 51 Packages Projects Releases Wiki Activity

[Distributed] deleteKey support for HashStore (#46049)

Browse Source 
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
This commit is contained in:
Omkar Salpekar
2020-10-14 11:56:05 -07:00
committed by Facebook GitHub Bot
parent 74f13a8b8f
commit 2ffb768607
2 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
torch/lib/c10d/HashStore.cpp
View File

@@ -84,8 +84,10 @@ int64_t HashStore::getNumKeys() {
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) {

7
torch/lib/c10d/test/HashStoreTest.cpp
View File

@@ -21,6 +21,13 @@ void testGetSet(std::string prefix = "") {
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_.