1 /** 2 * Copyright 2022-2023 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_CORE_MINDRT_RUNTIME_PARALLEL_THREAD_POOL_MANAGER_H_ 18 #define MINDSPORE_CORE_MINDRT_RUNTIME_PARALLEL_THREAD_POOL_MANAGER_H_ 19 20 #include <queue> 21 #include <vector> 22 #include <mutex> 23 #include <shared_mutex> 24 #include <unordered_map> 25 #include <condition_variable> 26 #include <map> 27 #include <string> 28 #include "thread/parallel_threadpool.h" 29 30 namespace mindspore { 31 class ThreadPool; 32 class ParallelThreadPool; 33 struct ParallelTask; 34 class ParallelWorker; 35 36 class MS_CORE_API ParallelThreadPoolManager { 37 public: 38 static ParallelThreadPoolManager *GetInstance(); 39 40 ~ParallelThreadPoolManager(); 41 42 void Init(bool enable_shared_thread_pool, const std::string &runner_id, int worker_num, int remaining_thread_num, 43 int task_num); 44 45 void SetHasIdlePool(std::string runner_id, bool is_idle); 46 47 void ResetParallelThreadPoolManager(const std::string &runner_id); 48 49 bool GetEnableSharedThreadPool(std::string runner_id); 50 51 void ActivatePool(const std::string &runner_id, int model_id); 52 53 void SetFreePool(const std::string &runner_id, int model_id); 54 55 void BindPoolToRunner(ThreadPool *pool, const std::map<std::string, std::map<std::string, std::string>> *config_info); 56 57 int GetThreadPoolSize(ThreadPool *pool); 58 59 int GetTaskNum(const std::map<std::string, std::map<std::string, std::string>> *config_info); 60 61 ParallelThreadPool *GetIdleThreadPool(const std::string &runner_id, ParallelTask *task); 62 63 private: 64 ParallelThreadPoolManager() = default; 65 66 private: 67 // runner id <=> thread pool(a model has a thread pool) 68 std::map<std::string, std::vector<ParallelThreadPool *>> runner_id_pools_; 69 // pool sorted by model worker id 70 std::unordered_map<ParallelThreadPool *, std::vector<ParallelWorker *>> pool_workers_; 71 72 std::shared_mutex pool_manager_mutex_; 73 std::map<std::string, bool> has_idle_pool_; 74 std::map<std::string, bool> enable_shared_thread_pool_; 75 std::map<std::string, int> runner_worker_num_; 76 std::map<std::string, int> worker_init_num_; 77 std::map<std::string, int> idle_pool_num_; 78 std::map<std::string, int> remaining_thread_num_; 79 std::map<std::string, int> thread_num_limit_; 80 }; 81 } // namespace mindspore 82 #endif // MINDSPORE_CORE_MINDRT_RUNTIME_PARALLEL_THREAD_POOL_MANAGER_H_ 83