forked from huawei/mindspore2022
117 lines
4.1 KiB
C++
117 lines
4.1 KiB
C++
/**
|
|
* Copyright 2021 Huawei Technologies Co., Ltd
|
|
*
|
|
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
* you may not use this file except in compliance with the License.
|
|
* You may obtain a copy of the License at
|
|
*
|
|
* http://www.apache.org/licenses/LICENSE-2.0
|
|
*
|
|
* Unless required by applicable law or agreed to in writing, software
|
|
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
* See the License for the specific language governing permissions and
|
|
* limitations under the License.
|
|
*/
|
|
|
|
#ifndef MIINDSPORE_CCSRC_DISTRIBUTED_PERSISTENT_DATA_H_
|
|
#define MIINDSPORE_CCSRC_DISTRIBUTED_PERSISTENT_DATA_H_
|
|
|
|
#include <map>
|
|
#include <memory>
|
|
#include <vector>
|
|
#include <string>
|
|
#include <thread>
|
|
#include <utility>
|
|
|
|
#include "distributed/persistent/storage/local_file.h"
|
|
#include "utils/log_adapter.h"
|
|
|
|
namespace mindspore {
|
|
namespace distributed {
|
|
namespace persistent {
|
|
// The data class is used to save and manage the tensor in memory, and provides
|
|
// interfaces for persistence and disaster recovery.
|
|
template <typename T>
|
|
class Data {
|
|
public:
|
|
explicit Data(const std::shared_ptr<std::vector<T>> &data, const std::shared_ptr<std::vector<int>> &shape = nullptr)
|
|
: data_(data), shape_(shape) {}
|
|
|
|
virtual ~Data() = default;
|
|
|
|
// Get the memory data of Data
|
|
T *data() const { return data_->data(); }
|
|
|
|
// Get the mutable memory data of Data
|
|
std::shared_ptr<std::vector<T>> MutableData() const { return data_; }
|
|
|
|
// Get the element number of Data
|
|
size_t size() const { return data_->size(); }
|
|
|
|
// Get the dimension information of Data.
|
|
std::shared_ptr<std::vector<int>> shape() const { return shape_; }
|
|
|
|
protected:
|
|
// Container used to store continuous memory buffer of Data.
|
|
std::shared_ptr<std::vector<T>> data_;
|
|
|
|
// Container used to record the dimension information of Data which persists a tensor.
|
|
std::shared_ptr<std::vector<int>> shape_;
|
|
};
|
|
|
|
// Implementation of the class Data to complete the function of persistence and disaster tolerance.
|
|
template <typename T>
|
|
class PersistentData : public Data<T> {
|
|
public:
|
|
explicit PersistentData(const std::shared_ptr<std::vector<T>> &data,
|
|
const std::shared_ptr<std::vector<int>> &shape = nullptr)
|
|
: Data<T>(data, shape) {}
|
|
|
|
~PersistentData() override = default;
|
|
|
|
// Initialize storage module.
|
|
// Custom storage config, you can choose different configurations according to different storage forms,
|
|
// such as using file storage by configuring the file storage path,
|
|
// and config can be like this: std::map<std::string, std::string> config = {{kFileStoragePath, "real_path_of_dir"}};
|
|
void Initialize(const std::map<std::string, std::string> &storage_config);
|
|
|
|
// In disaster recovery mode, memory of tensor need to be saved into disk file periodically.
|
|
void Persist(const storage::DirtyInfo &dirty_info) const;
|
|
|
|
// In disaster recovery mode, server node or worker node need to restore persistent data when restart.
|
|
void Restore() const;
|
|
|
|
private:
|
|
// The following variables are used in disaster recovery mode:
|
|
// The threads used to execute persistence task.
|
|
std::thread persist_thread_;
|
|
|
|
// The file storage handle used to persist data.
|
|
std::shared_ptr<storage::StorageBase> storage_;
|
|
};
|
|
|
|
template <typename T>
|
|
void PersistentData<T>::Initialize(const std::map<std::string, std::string> &storage_config) {
|
|
storage_ = std::make_shared<storage::LocalFile>(storage_config);
|
|
}
|
|
|
|
template <typename T>
|
|
void PersistentData<T>::Persist(const storage::DirtyInfo &dirty_info) const {
|
|
MS_EXCEPTION_IF_NULL(storage_);
|
|
storage::InputData input = std::make_tuple(*Data<T>::shape_, Data<T>::data(), Data<T>::size() * sizeof(T));
|
|
storage_->Write(input, dirty_info);
|
|
}
|
|
|
|
template <typename T>
|
|
void PersistentData<T>::Restore() const {
|
|
storage::OutputData output = std::make_pair(Data<T>::data(), Data<T>::size() * sizeof(T));
|
|
MS_EXCEPTION_IF_NULL(storage_);
|
|
storage_->Read(output);
|
|
}
|
|
} // namespace persistent
|
|
} // namespace distributed
|
|
} // namespace mindspore
|
|
|
|
#endif // MIINDSPORE_CCSRC_DISTRIBUTED_PERSISTENT_DATA_H_
|