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
10namespace c10d {
11
12class 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