mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
71ca600af9
Summary: Since caffe2 and torch have been consolidated, CAFFE2_API should be merged with TORCH_API. Addresses a TODO. Manually edited some references of the removed `CAFFE2_API`: * `CONTRIBUTING.md` * `caffe2/proto/CMakeLists.txt` * `cmake/ProtoBuf.cmake` * `c10/macros/Export.h` * `torch/csrc/WindowsTorchApiMacro.h` Pull Request resolved: https://github.com/pytorch/pytorch/pull/49496 Reviewed By: malfet, samestep Differential Revision: D25600726 Pulled By: janeyx99 fbshipit-source-id: 7e068d959e397ac183c097d7e9a9afeca5ddd782
91 lines
2.5 KiB
C++
91 lines
2.5 KiB
C++
#pragma once
|
|
|
|
#include <chrono>
|
|
#include <cstdint>
|
|
#include <stdexcept>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
#include "caffe2/core/common.h"
|
|
|
|
namespace caffe2 {
|
|
|
|
class TORCH_API StoreHandler {
|
|
public:
|
|
static constexpr std::chrono::milliseconds kDefaultTimeout =
|
|
std::chrono::seconds(30);
|
|
static constexpr std::chrono::milliseconds kNoTimeout =
|
|
std::chrono::milliseconds::zero();
|
|
|
|
virtual ~StoreHandler();
|
|
|
|
/*
|
|
* Set data for the key if it doesn't exist.
|
|
* If the key exists the data should be the same as the existing key.
|
|
*/
|
|
virtual void set(const std::string& name, const std::string& data) = 0;
|
|
|
|
/*
|
|
* Get the data for the key.
|
|
* The call should wait until the key is stored with specified timeout
|
|
* and return data if set else fail.
|
|
*/
|
|
virtual std::string get(
|
|
const std::string& name,
|
|
const std::chrono::milliseconds& timeout = kDefaultTimeout) = 0;
|
|
|
|
/*
|
|
* Does an atomic add operation on the key and returns the latest updated
|
|
* value.
|
|
* Note: To access the current value for this counter call with value = 0
|
|
*/
|
|
virtual int64_t add(const std::string& name, int64_t value) = 0;
|
|
|
|
/*
|
|
* Returns the number of keys in this store.
|
|
*/
|
|
virtual int64_t getNumKeys() = 0;
|
|
|
|
/*
|
|
* Removes the specified key from the store.
|
|
*/
|
|
virtual bool deleteKey(const std::string& key) = 0;
|
|
|
|
/*
|
|
* Check if a keys exist in the store.
|
|
*/
|
|
virtual bool check(const std::vector<std::string>& names) = 0;
|
|
|
|
/*
|
|
* Wait for Keys to be stored.
|
|
*/
|
|
virtual void wait(
|
|
const std::vector<std::string>& names,
|
|
const std::chrono::milliseconds& timeout = kDefaultTimeout) = 0;
|
|
};
|
|
|
|
/*
|
|
* The backing store is no longer available. It may have been deleted.
|
|
*/
|
|
struct TORCH_API StoreHandlerNotAvailableException
|
|
: public std::runtime_error {
|
|
explicit StoreHandlerNotAvailableException(const std::string& msg)
|
|
: std::runtime_error(msg) {}
|
|
};
|
|
|
|
#define STORE_HANDLER_NOT_AVAILABLE(...) \
|
|
throw ::caffe2::StoreHandlerNotAvailableException( \
|
|
::c10::str("[", __FILE__, ":", __LINE__, "] ", __VA_ARGS__));
|
|
|
|
/*
|
|
* Timeout accessing the store.
|
|
*/
|
|
struct TORCH_API StoreHandlerTimeoutException : public std::runtime_error {
|
|
explicit StoreHandlerTimeoutException(const std::string& msg)
|
|
: std::runtime_error(msg) {}
|
|
};
|
|
|
|
#define STORE_HANDLER_TIMEOUT(...) \
|
|
throw ::caffe2::StoreHandlerTimeoutException( \
|
|
::c10::str("[", __FILE__, ":", __LINE__, "] ", __VA_ARGS__));
|
|
} // namespace caffe2
|