mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
ec195129ec
Summary: This will allow users to set customized timeout option for the store. Tested by my own debug print to make sure that C++ actually used the timeout Pull Request resolved: https://github.com/pytorch/pytorch/pull/11265 Differential Revision: D9666164 Pulled By: teng-li fbshipit-source-id: 4eb6441783da106a3fd59b95457e503e83e4640f
36 lines
843 B
C++
36 lines
843 B
C++
#pragma once
|
|
|
|
#include <c10d/Store.hpp>
|
|
|
|
namespace c10d {
|
|
|
|
class PrefixStore : public Store {
|
|
public:
|
|
explicit PrefixStore(const std::string& prefix, Store& store);
|
|
|
|
virtual ~PrefixStore(){};
|
|
|
|
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:
|
|
std::string prefix_;
|
|
Store& store_;
|
|
|
|
std::string joinKey(const std::string& key);
|
|
std::vector<std::string> joinKeys(const std::vector<std::string>& keys);
|
|
};
|
|
|
|
} // namespace c10d
|