1 | |
2 | #include <c10/util/irange.h> |
3 | #include "StoreTestCommon.hpp" |
4 | |
5 | #ifndef _WIN32 |
6 | #include <unistd.h> |
7 | #endif |
8 | |
9 | #include <iostream> |
10 | #include <thread> |
11 | |
12 | #include <gtest/gtest.h> |
13 | |
14 | #include <torch/csrc/distributed/c10d/FileStore.hpp> |
15 | #include <torch/csrc/distributed/c10d/PrefixStore.hpp> |
16 | |
17 | #ifdef _WIN32 |
18 | std::string tmppath() { |
19 | return c10d::test::autoGenerateTmpFilePath(); |
20 | } |
21 | #else |
22 | std::string tmppath() { |
23 | const char* tmpdir = getenv("TMPDIR" ); |
24 | if (tmpdir == nullptr) { |
25 | tmpdir = "/tmp" ; |
26 | } |
27 | |
28 | // Create template |
29 | std::vector<char> tmp(256); |
30 | auto len = snprintf(tmp.data(), tmp.size(), "%s/testXXXXXX" , tmpdir); |
31 | tmp.resize(len); |
32 | |
33 | // Create temporary file |
34 | auto fd = mkstemp(&tmp[0]); |
35 | if (fd == -1) { |
36 | throw std::system_error(errno, std::system_category()); |
37 | } |
38 | close(fd); |
39 | return std::string(tmp.data(), tmp.size()); |
40 | } |
41 | #endif |
42 | |
43 | void testGetSet(std::string path, std::string prefix = "" ) { |
44 | // Basic Set/Get on File Store |
45 | { |
46 | auto fileStore = c10::make_intrusive<c10d::FileStore>(path, 2); |
47 | c10d::PrefixStore store(prefix, fileStore); |
48 | c10d::test::set(store, "key0" , "value0" ); |
49 | c10d::test::set(store, "key1" , "value1" ); |
50 | c10d::test::set(store, "key2" , "value2" ); |
51 | c10d::test::check(store, "key0" , "value0" ); |
52 | c10d::test::check(store, "key1" , "value1" ); |
53 | c10d::test::check(store, "key2" , "value2" ); |
54 | auto numKeys = fileStore->getNumKeys(); |
55 | EXPECT_EQ(numKeys, 4); |
56 | |
57 | // Check compareSet, does not check return value |
58 | c10d::test::compareSet(store, "key0" , "wrongExpectedValue" , "newValue" ); |
59 | c10d::test::check(store, "key0" , "value0" ); |
60 | c10d::test::compareSet(store, "key0" , "value0" , "newValue" ); |
61 | c10d::test::check(store, "key0" , "newValue" ); |
62 | |
63 | // Check deleteKey |
64 | c10d::test::deleteKey(store, "key1" ); |
65 | numKeys = fileStore->getNumKeys(); |
66 | EXPECT_EQ(numKeys, 3); |
67 | c10d::test::check(store, "key0" , "newValue" ); |
68 | c10d::test::check(store, "key2" , "value2" ); |
69 | |
70 | c10d::test::set(store, "-key0" , "value-" ); |
71 | c10d::test::check(store, "key0" , "newValue" ); |
72 | c10d::test::check(store, "-key0" , "value-" ); |
73 | numKeys = fileStore->getNumKeys(); |
74 | EXPECT_EQ(numKeys, 4); |
75 | c10d::test::deleteKey(store, "-key0" ); |
76 | numKeys = fileStore->getNumKeys(); |
77 | EXPECT_EQ(numKeys, 3); |
78 | c10d::test::check(store, "key0" , "newValue" ); |
79 | c10d::test::check(store, "key2" , "value2" ); |
80 | } |
81 | |
82 | // Perform get on new instance |
83 | { |
84 | auto fileStore = c10::make_intrusive<c10d::FileStore>(path, 2); |
85 | c10d::PrefixStore store(prefix, fileStore); |
86 | c10d::test::check(store, "key0" , "newValue" ); |
87 | auto numKeys = fileStore->getNumKeys(); |
88 | // There will be 4 keys since we still use the same underlying file as the |
89 | // other store above. |
90 | EXPECT_EQ(numKeys, 4); |
91 | } |
92 | } |
93 | |
94 | void stressTestStore(std::string path, std::string prefix = "" ) { |
95 | // Hammer on FileStore::add |
96 | const auto numThreads = 4; |
97 | const auto numIterations = 100; |
98 | |
99 | std::vector<std::thread> threads; |
100 | c10d::test::Semaphore sem1, sem2; |
101 | |
102 | for (C10_UNUSED const auto i : c10::irange(numThreads)) { |
103 | threads.emplace_back(std::thread([&] { |
104 | auto fileStore = |
105 | c10::make_intrusive<c10d::FileStore>(path, numThreads + 1); |
106 | c10d::PrefixStore store(prefix, fileStore); |
107 | sem1.post(); |
108 | sem2.wait(); |
109 | for (C10_UNUSED const auto j : c10::irange(numIterations)) { |
110 | store.add("counter" , 1); |
111 | } |
112 | })); |
113 | } |
114 | |
115 | sem1.wait(numThreads); |
116 | sem2.post(numThreads); |
117 | for (auto& thread : threads) { |
118 | thread.join(); |
119 | } |
120 | |
121 | // Check that the counter has the expected value |
122 | { |
123 | auto fileStore = c10::make_intrusive<c10d::FileStore>(path, numThreads + 1); |
124 | c10d::PrefixStore store(prefix, fileStore); |
125 | std::string expected = std::to_string(numThreads * numIterations); |
126 | c10d::test::check(store, "counter" , expected); |
127 | } |
128 | } |
129 | |
130 | class FileStoreTest : public ::testing::Test { |
131 | protected: |
132 | void SetUp() override { |
133 | path_ = tmppath(); |
134 | } |
135 | |
136 | void TearDown() override { |
137 | unlink(path_.c_str()); |
138 | } |
139 | |
140 | std::string path_; |
141 | }; |
142 | |
143 | TEST_F(FileStoreTest, testGetAndSet) { |
144 | testGetSet(path_); |
145 | } |
146 | |
147 | TEST_F(FileStoreTest, testGetAndSetWithPrefix) { |
148 | testGetSet(path_, "testPrefix" ); |
149 | } |
150 | |
151 | TEST_F(FileStoreTest, testStressStore) { |
152 | stressTestStore(path_); |
153 | } |
154 | |
155 | TEST_F(FileStoreTest, testStressStoreWithPrefix) { |
156 | stressTestStore(path_, "testPrefix" ); |
157 | } |
158 | |