mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
97036d3c30
Summary: This addressed: https://github.com/pytorch/pytorch/issues/11874 and we will have the identical file init_method behavior as the previous THD file init. Also the FileStore::add bug is pretty annoying. Two bugs: (1) Add doesn't append to the end of the file. (2) Cache doesn't get updated. Both are fixed and tests are covered. I examined the /tmp to ensure that all temp files are auto deleted after test_c10d.py Pull Request resolved: https://github.com/pytorch/pytorch/pull/13708 Reviewed By: pietern Differential Revision: D12972810 Pulled By: teng-li fbshipit-source-id: 917255390aa52845f6b0ad0f283875a7a704da48
45 lines
968 B
C++
45 lines
968 B
C++
#pragma once
|
|
|
|
#include <sys/types.h>
|
|
|
|
#include <unordered_map>
|
|
|
|
#include <c10d/Store.hpp>
|
|
|
|
namespace c10d {
|
|
|
|
class FileStore : public Store {
|
|
public:
|
|
explicit FileStore(const std::string& path, int numWorkers);
|
|
|
|
virtual ~FileStore();
|
|
|
|
void set(const std::string& key, const std::vector<uint8_t>& value) override;
|
|
|
|
std::vector<uint8_t> get(const std::string& key) override;
|
|
|
|
int64_t add(const std::string& key, int64_t value) override;
|
|
|
|
bool check(const std::vector<std::string>& keys) override;
|
|
|
|
void wait(const std::vector<std::string>& keys) override;
|
|
|
|
void wait(
|
|
const std::vector<std::string>& keys,
|
|
const std::chrono::milliseconds& timeout) override;
|
|
|
|
protected:
|
|
int64_t addHelper(const std::string& key, int64_t i);
|
|
|
|
std::string path_;
|
|
off_t pos_;
|
|
|
|
int numWorkers_;
|
|
const std::string cleanupKey_;
|
|
const std::string regularPrefix_;
|
|
|
|
std::unordered_map<std::string, std::vector<uint8_t>> cache_;
|
|
};
|
|
|
|
} // namespace c10d
|