Move processor interface init from learner to communicator functional

ZiyueXu77 · Apr 12, 2024 · 30b7ed5 · 30b7ed5
1 parent 1107604
commit 30b7ed5
Showing 1 changed file with 19 additions and 7 deletions.
diff --git a/src/collective/communicator.cc b/src/collective/communicator.cc
@@ -41,13 +41,25 @@ void Communicator::Init(Json const& config) {
     case CommunicatorType::kFederated: {
 #if defined(XGBOOST_USE_FEDERATED)
       communicator_.reset(FederatedCommunicator::Create(config));
-      std::cout << "!!!!!!!! Communicator Initialization!!!!!!!!!!!!!!!!!!!! " << std::endl;
-      auto plugin_name = "dummy";
-      std::map<std::string, std::string> loader_params = {{"LIBRARY_PATH", "/tmp"}};
-      std::map<std::string, std::string> proc_params = {};
-      processing::ProcessorLoader loader(loader_params);
-      processor_instance = loader.load(plugin_name);
-      processor_instance->Initialize(collective::GetRank() == 0, proc_params);
+      // Get processor configs
+      std::string plugin_name{};
+      std::string loader_params_key{};
+      std::string loader_params_map{};
+      std::string proc_params_key{};
+      std::string proc_params_map{};
+      plugin_name = OptionalArg<String>(config, "plugin_name", plugin_name);
+      loader_params_key = OptionalArg<String>(config, "loader_params_key", loader_params_key);
+      loader_params_map = OptionalArg<String>(config, "loader_params_map", loader_params_map);
+      proc_params_key = OptionalArg<String>(config, "proc_params_key", proc_params_key);
+      proc_params_map = OptionalArg<String>(config, "proc_params_map", proc_params_map);
+      // Initialize processor if plugin_name is provided
+      if (!plugin_name.empty()){
+        std::map<std::string, std::string> loader_params = {{loader_params_key, loader_params_map}};
+        std::map<std::string, std::string> proc_params = {{proc_params_key, proc_params_map}};
+        processing::ProcessorLoader loader(loader_params);
+        processor_instance = loader.load(plugin_name);
+        processor_instance->Initialize(collective::GetRank() == 0, proc_params);
+      }
 #else
       LOG(FATAL) << "XGBoost is not compiled with Federated Learning support.";
 #endif