1 | #include <torch/csrc/distributed/c10d/PrefixStore.hpp> |
2 | #include <utility> |
3 | |
4 | namespace c10d { |
5 | |
6 | PrefixStore::PrefixStore(std::string prefix, c10::intrusive_ptr<Store> store) |
7 | : prefix_(std::move(prefix)), store_(std::move(store)) {} |
8 | |
9 | std::string PrefixStore::joinKey(const std::string& key) { |
10 | return prefix_ + "/" + key; |
11 | } |
12 | |
13 | std::vector<std::string> PrefixStore::joinKeys( |
14 | const std::vector<std::string>& keys) { |
15 | std::vector<std::string> joinedKeys; |
16 | joinedKeys.reserve(keys.size()); |
17 | for (const auto& key : keys) { |
18 | joinedKeys.emplace_back(joinKey(key)); |
19 | } |
20 | return joinedKeys; |
21 | } |
22 | |
23 | void PrefixStore::set( |
24 | const std::string& key, |
25 | const std::vector<uint8_t>& value) { |
26 | store_->set(joinKey(key), value); |
27 | } |
28 | |
29 | std::vector<uint8_t> PrefixStore::compareSet( |
30 | const std::string& key, |
31 | const std::vector<uint8_t>& expectedValue, |
32 | const std::vector<uint8_t>& desiredValue) { |
33 | return store_->compareSet(joinKey(key), expectedValue, desiredValue); |
34 | } |
35 | |
36 | std::vector<uint8_t> PrefixStore::get(const std::string& key) { |
37 | return store_->get(joinKey(key)); |
38 | } |
39 | |
40 | int64_t PrefixStore::add(const std::string& key, int64_t value) { |
41 | return store_->add(joinKey(key), value); |
42 | } |
43 | |
44 | bool PrefixStore::deleteKey(const std::string& key) { |
45 | return store_->deleteKey(joinKey(key)); |
46 | } |
47 | |
48 | void PrefixStore::watchKey(const std::string& key, WatchKeyCallback callback) { |
49 | return store_->watchKey(joinKey(key), std::move(callback)); |
50 | } |
51 | |
52 | int64_t PrefixStore::getNumKeys() { |
53 | return store_->getNumKeys(); |
54 | } |
55 | |
56 | bool PrefixStore::check(const std::vector<std::string>& keys) { |
57 | auto joinedKeys = joinKeys(keys); |
58 | return store_->check(joinedKeys); |
59 | } |
60 | |
61 | void PrefixStore::wait(const std::vector<std::string>& keys) { |
62 | auto joinedKeys = joinKeys(keys); |
63 | store_->wait(joinedKeys); |
64 | } |
65 | |
66 | void PrefixStore::wait( |
67 | const std::vector<std::string>& keys, |
68 | const std::chrono::milliseconds& timeout) { |
69 | auto joinedKeys = joinKeys(keys); |
70 | store_->wait(joinedKeys, timeout); |
71 | } |
72 | |
73 | const std::chrono::milliseconds& PrefixStore::getTimeout() const noexcept { |
74 | return store_->getTimeout(); |
75 | } |
76 | |
77 | void PrefixStore::setTimeout(const std::chrono::milliseconds& timeout) { |
78 | store_->setTimeout(timeout); |
79 | } |
80 | |
81 | c10::intrusive_ptr<Store> PrefixStore::getUnderlyingStore() { |
82 | return store_; |
83 | } |
84 | |
85 | } // namespace c10d |
86 | |