1 /** 2 * Copyright 2021 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_FL_SERVER_MODEL_STORE_H_ 18 #define MINDSPORE_CCSRC_FL_SERVER_MODEL_STORE_H_ 19 20 #include <map> 21 #include <memory> 22 #include <string> 23 #include "fl/server/common.h" 24 #include "fl/server/memory_register.h" 25 #include "fl/server/executor.h" 26 27 namespace mindspore { 28 namespace fl { 29 namespace server { 30 // The initial iteration number is 0 in server. 31 constexpr size_t kInitIterationNum = 0; 32 33 // The initial iteration number after ModelStore is reset. 34 constexpr size_t kResetInitIterNum = 1; 35 36 // Server framework use ModelStore to store and query models. 37 // ModelStore stores multiple models because worker could get models of the previous iterations. 38 class ModelStore { 39 public: GetInstance()40 static ModelStore &GetInstance() { 41 static ModelStore instance; 42 return instance; 43 } 44 45 // Initialize ModelStore with max count of models need to be stored. 46 void Initialize(uint32_t max_count = 3); 47 48 // Store the model of the given iteration. The model is acquired from Executor. If the current model count is already 49 // max_model_count_, the earliest model will be replaced. 50 void StoreModelByIterNum(size_t iteration, const std::map<std::string, AddressPtr> &model); 51 52 // Get model of the given iteration. 53 std::map<std::string, AddressPtr> GetModelByIterNum(size_t iteration); 54 55 // Reset the stored models. Called when federated learning job finishes. 56 void Reset(); 57 58 // Returns all models stored in ModelStore. 59 const std::map<size_t, std::shared_ptr<MemoryRegister>> &iteration_to_model(); 60 61 // Returns the model size, which could be calculated at the initializing phase. 62 size_t model_size() const; 63 64 private: ModelStore()65 ModelStore() : max_model_count_(0), model_size_(0), iteration_to_model_({}) {} 66 ~ModelStore() = default; 67 ModelStore(const ModelStore &) = delete; 68 ModelStore &operator=(const ModelStore &) = delete; 69 70 // To store multiple models, new memory must assigned. The max memory size assigned for models is max_model_count_ * 71 // model_size_. 72 std::shared_ptr<MemoryRegister> AssignNewModelMemory(); 73 74 // Calculate the model size. This method should be called after iteration_to_model_ is initialized. 75 size_t ComputeModelSize(); 76 77 size_t max_model_count_; 78 size_t model_size_; 79 80 // Initial model which is the model of iteration 0. 81 std::shared_ptr<MemoryRegister> initial_model_; 82 83 // The number of all models stored is max_model_count_. 84 std::mutex model_mtx_; 85 std::map<size_t, std::shared_ptr<MemoryRegister>> iteration_to_model_; 86 }; 87 } // namespace server 88 } // namespace fl 89 } // namespace mindspore 90 #endif // MINDSPORE_CCSRC_FL_SERVER_MODEL_STORE_H_ 91