1 | #pragma once |
---|---|
2 | |
3 | #include <sys/types.h> |
4 | |
5 | #include <mutex> |
6 | #include <unordered_map> |
7 | |
8 | #include <torch/csrc/distributed/c10d/Store.hpp> |
9 | |
10 | namespace c10d { |
11 | |
12 | class TORCH_API FileStore : public Store { |
13 | public: |
14 | explicit FileStore(std::string path, int numWorkers); |
15 | |
16 | ~FileStore() override; |
17 | |
18 | void set(const std::string& key, const std::vector<uint8_t>& value) override; |
19 | |
20 | std::vector<uint8_t> compareSet( |
21 | const std::string& key, |
22 | const std::vector<uint8_t>& expectedValue, |
23 | const std::vector<uint8_t>& desiredValue) override; |
24 | |
25 | std::vector<uint8_t> get(const std::string& key) override; |
26 | |
27 | int64_t add(const std::string& key, int64_t value) override; |
28 | |
29 | int64_t getNumKeys() override; |
30 | |
31 | bool deleteKey(const std::string& key) override; |
32 | |
33 | bool check(const std::vector<std::string>& keys) override; |
34 | |
35 | void wait(const std::vector<std::string>& keys) override; |
36 | |
37 | void wait( |
38 | const std::vector<std::string>& keys, |
39 | const std::chrono::milliseconds& timeout) override; |
40 | |
41 | // Returns the path used by the FileStore. |
42 | const std::string& getPath() const noexcept { |
43 | return path_; |
44 | } |
45 | |
46 | protected: |
47 | int64_t addHelper(const std::string& key, int64_t i); |
48 | |
49 | std::string path_; |
50 | off_t pos_{0}; |
51 | |
52 | int numWorkers_; |
53 | const std::string cleanupKey_; |
54 | const std::string refCountKey_; |
55 | const std::string regularPrefix_; |
56 | const std::string deletePrefix_; |
57 | |
58 | std::unordered_map<std::string, std::vector<uint8_t>> cache_; |
59 | |
60 | std::mutex activeFileOpLock_; |
61 | }; |
62 | |
63 | } // namespace c10d |
64 |