mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
12f97bb2e9
Fixes the string_view errors and reland the work. The previous changes in torch/csrc/utils/invalid_arguments.cpp were too aggressive and not tested thoroughly. They are discarded. Pull Request resolved: https://github.com/pytorch/pytorch/pull/110518 Approved by: https://github.com/ezyang
102 lines
2.8 KiB
C++
102 lines
2.8 KiB
C++
#pragma once
|
|
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include <c10/macros/Macros.h>
|
|
#include <torch/custom_class.h>
|
|
|
|
namespace c10d {
|
|
|
|
// callback function will be given arguments (optional<string> oldValue,
|
|
// optional<string> newValue)
|
|
using WatchKeyCallback =
|
|
std::function<void(c10::optional<std::string>, c10::optional<std::string>)>;
|
|
|
|
class TORCH_API Store : public torch::CustomClassHolder {
|
|
public:
|
|
static constexpr std::chrono::milliseconds kDefaultTimeout =
|
|
std::chrono::seconds(300);
|
|
static constexpr std::chrono::milliseconds kNoTimeout =
|
|
std::chrono::milliseconds::zero();
|
|
|
|
Store() : timeout_(kDefaultTimeout) {}
|
|
|
|
explicit Store(const std::chrono::milliseconds& timeout)
|
|
: timeout_(timeout) {}
|
|
|
|
Store(const Store&) = default;
|
|
Store(Store&&) noexcept = default;
|
|
|
|
~Store() override = default;
|
|
|
|
void set(const std::string& key, const std::string& value);
|
|
|
|
virtual void set(
|
|
const std::string& key,
|
|
const std::vector<uint8_t>& value) = 0;
|
|
|
|
std::string compareSet(
|
|
const std::string& key,
|
|
const std::string& currentValue,
|
|
const std::string& newValue);
|
|
|
|
virtual std::vector<uint8_t> compareSet(
|
|
const std::string& key,
|
|
const std::vector<uint8_t>& currentValue,
|
|
const std::vector<uint8_t>& newValue) {
|
|
TORCH_INTERNAL_ASSERT(false, "Not implemented.");
|
|
}
|
|
|
|
std::string get_to_str(const std::string& key);
|
|
|
|
virtual std::vector<uint8_t> get(const std::string& key) = 0;
|
|
|
|
virtual int64_t add(const std::string& key, int64_t value) = 0;
|
|
|
|
virtual bool deleteKey(const std::string& key) = 0;
|
|
|
|
virtual bool check(const std::vector<std::string>& keys) = 0;
|
|
|
|
virtual int64_t getNumKeys() = 0;
|
|
|
|
virtual void wait(const std::vector<std::string>& keys) = 0;
|
|
|
|
virtual void wait(
|
|
const std::vector<std::string>& keys,
|
|
const std::chrono::milliseconds& timeout) = 0;
|
|
|
|
virtual const std::chrono::milliseconds& getTimeout() const noexcept;
|
|
|
|
virtual void setTimeout(const std::chrono::milliseconds& timeout);
|
|
|
|
// watchKey() is deprecated and no longer supported.
|
|
virtual void watchKey(
|
|
const std::string& /* unused */,
|
|
WatchKeyCallback /* unused */) {
|
|
TORCH_CHECK(false, "watchKey is deprecated, no implementation support it.");
|
|
}
|
|
|
|
virtual void append(
|
|
const std::string& key,
|
|
const std::vector<uint8_t>& value);
|
|
|
|
virtual std::vector<std::vector<uint8_t>> multiGet(
|
|
const std::vector<std::string>& keys);
|
|
|
|
virtual void multiSet(
|
|
const std::vector<std::string>& keys,
|
|
const std::vector<std::vector<uint8_t>>& values);
|
|
|
|
// Returns true if this store support append, multiGet and multiSet
|
|
virtual bool hasExtendedApi() const;
|
|
|
|
protected:
|
|
std::chrono::milliseconds timeout_;
|
|
};
|
|
|
|
} // namespace c10d
|