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. getPath() const42 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