1 #include <torch/csrc/distributed/c10d/Store.hpp>
2
3 namespace c10d {
4
getTimeout() const5 const std::chrono::milliseconds& Store::getTimeout() const noexcept {
6 return timeout_;
7 }
8
9 // Set timeout function
setTimeout(const std::chrono::milliseconds & timeout)10 void Store::setTimeout(const std::chrono::milliseconds& timeout) {
11 timeout_ = timeout;
12 }
13
set(const std::string & key,const std::string & value)14 void Store::set(const std::string& key, const std::string& value) {
15 set(key, std::vector<uint8_t>(value.begin(), value.end()));
16 }
17
compareSet(const std::string & key,const std::string & currentValue,const std::string & newValue)18 std::string Store::compareSet(
19 const std::string& key,
20 const std::string& currentValue,
21 const std::string& newValue) {
22 auto value = compareSet(
23 key,
24 std::vector<uint8_t>(currentValue.begin(), currentValue.end()),
25 std::vector<uint8_t>(newValue.begin(), newValue.end()));
26 return std::string(value.begin(), value.end());
27 }
28
get_to_str(const std::string & key)29 std::string Store::get_to_str(const std::string& key) {
30 auto value = get(key);
31 return std::string(value.begin(), value.end());
32 }
33
append(const std::string & key,const std::vector<uint8_t> & value)34 void Store::append(const std::string& key, const std::vector<uint8_t>& value) {
35 // This fallback depends on compareSet
36 std::vector<uint8_t> expected = value;
37 std::vector<uint8_t> current;
38 // cannot use get(key) as it might block forever if the key doesn't exist
39 current = compareSet(key, current, expected);
40 while (current != expected) {
41 expected = current;
42 expected.insert(expected.end(), value.begin(), value.end());
43 current = compareSet(key, current, expected);
44 }
45 }
46
multiGet(const std::vector<std::string> & keys)47 std::vector<std::vector<uint8_t>> Store::multiGet(
48 const std::vector<std::string>& keys) {
49 std::vector<std::vector<uint8_t>> result;
50 result.reserve(keys.size());
51 for (auto& key : keys) {
52 result.emplace_back(get(key));
53 }
54 return result;
55 }
56
multiSet(const std::vector<std::string> & keys,const std::vector<std::vector<uint8_t>> & values)57 void Store::multiSet(
58 const std::vector<std::string>& keys,
59 const std::vector<std::vector<uint8_t>>& values) {
60 for (auto i : ::c10::irange(keys.size())) {
61 set(keys[i], values[i]);
62 }
63 }
64
hasExtendedApi() const65 bool Store::hasExtendedApi() const {
66 return false;
67 }
68
69 } // namespace c10d
70