diff --git a/mindspore/ccsrc/frontend/parallel/graph_util/graph_splitter.cc b/mindspore/ccsrc/frontend/parallel/graph_util/graph_splitter.cc index e4061a8d29a..e86cfa8cd6a 100644 --- a/mindspore/ccsrc/frontend/parallel/graph_util/graph_splitter.cc +++ b/mindspore/ccsrc/frontend/parallel/graph_util/graph_splitter.cc @@ -196,15 +196,26 @@ void GraphSplitter::Run() { MS_EXCEPTION_IF_NULL(func_graph_); MS_EXCEPTION_IF_NULL(func_graph_->manager()); + // Step 1: Dye all the nodes of the whole func_graph_. DyeGraph(); + // If all nodes are all on this process, no need to split the graph. So return. + if (std::find_if(node_labels_.begin(), node_labels_.end(), [&](const auto &node_to_label) { + return node_to_label.second != this_process_label_; + }) == node_labels_.end()) { + return; + } + // Step 2: Generate the node segments with different labels. std::vector segments = GenerateSplitSegments(); // If the segment number is 0, there will be no distributed execution. if (segments.empty()) { return; } + + // Step 3: Create inter-process operators for segments with different labels. InterProcessOpEdgesInfo comm_edges = GenerateInterProcessOperators(); + // Step 4: Split the graph and eliminate extra nodes. SplitGraph(segments, comm_edges); } diff --git a/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc b/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc index deadadab142..70128dee307 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc +++ b/mindspore/ccsrc/runtime/graph_scheduler/graph_scheduler.cc @@ -434,7 +434,6 @@ void GraphScheduler::Initialize() { // Create and initialize RpcNodeScheduler. rpc_node_scheduler_ = std::make_unique(); MS_EXCEPTION_IF_NULL(rpc_node_scheduler_); - rpc_node_scheduler_->Initialize(); #endif BuildAndScheduleGlobalActor(); @@ -544,7 +543,7 @@ void GraphScheduler::Schedule(const ActorSet *actor_set) { #ifdef ENABLE_RPC_ACTOR // Build physical connections in 'RpcNodeScheduler::Schedule()' method. This costs some time. MS_EXCEPTION_IF_NULL(rpc_node_scheduler_); - rpc_node_scheduler_->Schedule(); + rpc_node_scheduler_->Schedule(actor_set); #endif } @@ -565,7 +564,8 @@ void GraphScheduler::Run(ActorSet *const actor_set, const std::vector(rpc_node_scheduler_.get(), &op_context); + auto op_context_setter = + std::make_shared(rpc_node_scheduler_.get(), actor_set->rpc_actors_, &op_context); MS_EXCEPTION_IF_NULL(op_context_setter); #endif @@ -696,7 +696,7 @@ ActorSetPtr GraphScheduler::Build(const GraphCompilerInfo &graph_compiler_info) #ifdef ENABLE_RPC_ACTOR MS_EXCEPTION_IF_NULL(rpc_node_scheduler_); - actor_set->rpc_actors_ = rpc_node_scheduler_->Build(graph_compiler_info); + actor_set->rpc_actors_ = rpc_node_scheduler_->Build(actor_set.get()); #endif return actor_set; } @@ -1095,14 +1095,12 @@ KernelActorPtr GraphScheduler::GenerateRpcActor(const CNodePtr &kernel, const De std::make_shared(kernel->fullname_with_scope(), kernel, device_context, memory_manager_aid_, debug_aid_, recorder_aid_, strategy, ref_input_indexes, ref_output_indexes); MS_EXCEPTION_IF_NULL(send_actor); - rpc_node_scheduler_->InsertSendActor(send_actor); return send_actor; } else if (common::AnfAlgo::GetCNodeName(kernel) == kRpcRecvOpName) { auto recv_actor = std::make_shared(kernel->fullname_with_scope(), kernel, device_context, memory_manager_aid_, debug_aid_, recorder_aid_, strategy, ref_input_indexes, ref_output_indexes); MS_EXCEPTION_IF_NULL(recv_actor); - rpc_node_scheduler_->InsertRecvActor(recv_actor); return recv_actor; } else { MS_LOG(EXCEPTION) << "Kernel " << kernel->fullname_with_scope() << " is not an rpc kernel."; diff --git a/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.cc b/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.cc index 8bc6f65797a..c706dfe5f0e 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.cc +++ b/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.cc @@ -19,16 +19,30 @@ namespace mindspore { namespace runtime { -void RpcNodeScheduler::Initialize() { - rpc_actor_set_ = std::make_shared(); - MS_EXCEPTION_IF_NULL(rpc_actor_set_); -} +RpcActorSetPtr RpcNodeScheduler::Build(const ActorSet *actor_set) { + MS_EXCEPTION_IF_NULL(actor_set); + + // RpcActor inherits from KernelActor, so we need to filter out the rpc actors from kernel actors list. + std::vector kernel_actors = actor_set->kernel_actors_; + RpcActorSetPtr rpc_actor_set = std::make_shared(); + MS_EXCEPTION_IF_NULL(rpc_actor_set); -RpcActorSetPtr RpcNodeScheduler::Build(const GraphCompilerInfo &) { - MS_EXCEPTION_IF_NULL(rpc_actor_set_); std::vector rpc_actors; - (void)rpc_actors.insert(rpc_actors.end(), rpc_actor_set_->send_actors_.begin(), rpc_actor_set_->send_actors_.end()); - (void)rpc_actors.insert(rpc_actors.end(), rpc_actor_set_->recv_actors_.begin(), rpc_actor_set_->recv_actors_.end()); + for (const auto &kernel_actor : kernel_actors) { + auto rpc_actor = std::dynamic_pointer_cast(kernel_actor); + if (std::dynamic_pointer_cast(kernel_actor) == nullptr) { + continue; + } else { + rpc_actors.emplace_back(rpc_actor); + if (std::dynamic_pointer_cast(rpc_actor) != nullptr) { + rpc_actor_set->send_actors_.emplace_back(std::dynamic_pointer_cast(rpc_actor)); + } else if (std::dynamic_pointer_cast(rpc_actor) != nullptr) { + rpc_actor_set->recv_actors_.emplace_back(std::dynamic_pointer_cast(rpc_actor)); + } else { + MS_LOG(EXCEPTION) << "Rpc actor should be either SendActor or RecvActor."; + } + } + } // Create route table proxy for each rpc actor and set. for (auto &rpc_actor : rpc_actors) { @@ -37,13 +51,16 @@ RpcActorSetPtr RpcNodeScheduler::Build(const GraphCompilerInfo &) { rpc_actor->SetActorRouteRableProxy(proxy); } - return rpc_actor_set_; + return rpc_actor_set; } -void RpcNodeScheduler::Link(const ActorSet *) { - MS_EXCEPTION_IF_NULL(rpc_actor_set_); - std::vector send_actors = rpc_actor_set_->send_actors_; - std::vector recv_actors = rpc_actor_set_->recv_actors_; +void RpcNodeScheduler::Link(const ActorSet *actor_set) { + MS_EXCEPTION_IF_NULL(actor_set); + RpcActorSetPtr rpc_actor_set = actor_set->rpc_actors_; + MS_EXCEPTION_IF_NULL(rpc_actor_set); + std::vector send_actors = rpc_actor_set->send_actors_; + std::vector recv_actors = rpc_actor_set->recv_actors_; + // The inter-process edge is connected to a remote peer. So the peer info attributes in the kernel should be // sufficient for route table. for (auto &send_actor : send_actors) { @@ -84,19 +101,21 @@ void RpcNodeScheduler::Link(const ActorSet *) { } } -void RpcNodeScheduler::Schedule() { - MS_EXCEPTION_IF_NULL(rpc_actor_set_); +void RpcNodeScheduler::Schedule(const ActorSet *actor_set) { + MS_EXCEPTION_IF_NULL(actor_set); + RpcActorSetPtr rpc_actor_set = actor_set->rpc_actors_; + MS_EXCEPTION_IF_NULL(rpc_actor_set); // Must start server and register route table before looking up route and connecting. // Start servers of recv actors and register route table. - for (auto &recv_actor : rpc_actor_set_->recv_actors_) { + for (auto &recv_actor : rpc_actor_set->recv_actors_) { MS_EXCEPTION_IF_NULL(recv_actor); if (!recv_actor->StartServer()) { MS_LOG(EXCEPTION) << "Failed to start server for the recv actor."; } } // Lookup route and connect to servers for send actors. - for (auto &send_actor : rpc_actor_set_->send_actors_) { + for (auto &send_actor : rpc_actor_set->send_actors_) { MS_EXCEPTION_IF_NULL(send_actor); if (!send_actor->ConnectServer()) { MS_LOG(EXCEPTION) << "Failed to connect servers for the send actor."; @@ -104,40 +123,28 @@ void RpcNodeScheduler::Schedule() { } } -void RpcNodeScheduler::InsertSendActor(const SendActorPtr &send_actor) { - MS_EXCEPTION_IF_NULL(rpc_actor_set_); - MS_EXCEPTION_IF_NULL(send_actor); - (void)rpc_actor_set_->send_actors_.emplace_back(send_actor); -} - -void RpcNodeScheduler::InsertRecvActor(const RecvActorPtr &recv_actor) { - MS_EXCEPTION_IF_NULL(rpc_actor_set_); - MS_EXCEPTION_IF_NULL(recv_actor); - (void)rpc_actor_set_->recv_actors_.emplace_back(recv_actor); -} - -void RpcNodeScheduler::SetOpcontext(OpContext *const op_context) { +void RpcNodeScheduler::SetOpcontext(const RpcActorSetPtr &rpc_actors, OpContext *const op_context) { MS_EXCEPTION_IF_NULL(op_context); - MS_EXCEPTION_IF_NULL(rpc_actor_set_); + MS_EXCEPTION_IF_NULL(rpc_actors); - for (auto &recv_actor : rpc_actor_set_->recv_actors_) { + for (auto &recv_actor : rpc_actors->recv_actors_) { MS_EXCEPTION_IF_NULL(recv_actor); recv_actor->SetOpcontext(op_context); } - for (auto &send_actor : rpc_actor_set_->send_actors_) { + for (auto &send_actor : rpc_actors->send_actors_) { MS_EXCEPTION_IF_NULL(send_actor); send_actor->SetOpcontext(op_context); } } -void RpcNodeScheduler::ResetOpcontext() { - MS_EXCEPTION_IF_NULL(rpc_actor_set_); +void RpcNodeScheduler::ResetOpcontext(const RpcActorSetPtr &rpc_actors) { + MS_EXCEPTION_IF_NULL(rpc_actors); - for (auto &recv_actor : rpc_actor_set_->recv_actors_) { + for (auto &recv_actor : rpc_actors->recv_actors_) { MS_EXCEPTION_IF_NULL(recv_actor); recv_actor->ResetOpcontext(); } - for (auto &send_actor : rpc_actor_set_->send_actors_) { + for (auto &send_actor : rpc_actors->send_actors_) { MS_EXCEPTION_IF_NULL(send_actor); send_actor->ResetOpcontext(); } diff --git a/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.h b/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.h index d0fd75aa582..246ea768c5e 100644 --- a/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.h +++ b/mindspore/ccsrc/runtime/graph_scheduler/rpc_node_scheduler.h @@ -32,50 +32,43 @@ using mindspore::session::KernelWithIndex; // Scheduler for rpc actors, e.g., it adds inter-process arrows, generate router for actors, etc. class RpcNodeScheduler { public: - RpcNodeScheduler() : rpc_actor_set_(nullptr) {} + RpcNodeScheduler() = default; ~RpcNodeScheduler() = default; - // Create rpc actor set. - void Initialize(); - // Build rpc actors and return rpc actor set. - RpcActorSetPtr Build(const GraphCompilerInfo &graph_compiler_info); + RpcActorSetPtr Build(const ActorSet *actor_set); // Link rpc actors with inter-process arrows. void Link(const ActorSet *actor_set); // This should be called by 'GraphScheduler::Scheduler()' method. // Used to start servers for recv actors and create connections for send actors. - void Schedule(); - - // Insert Send/Recv actors generated by GraphScheduler to the rpc actor set. - void InsertSendActor(const SendActorPtr &send_actor); - void InsertRecvActor(const RecvActorPtr &recv_actor); + void Schedule(const ActorSet *actor_set); // Set op context to rpc actors. - void SetOpcontext(OpContext *const op_context); + void SetOpcontext(const RpcActorSetPtr &rpc_actors, OpContext *const op_context); // Reset op context for rpc actors. - void ResetOpcontext(); + void ResetOpcontext(const RpcActorSetPtr &rpc_actors); private: // Create new route table proxy. ActorRouteTableProxyPtr CreateRouteTableProxy(); - - RpcActorSetPtr rpc_actor_set_; }; // The setter of op context for rpc actors. class RpcActorOpContextSetter { public: - explicit RpcActorOpContextSetter(RpcNodeScheduler *rpc_node_scheduler, OpContext *const op_context) - : rpc_node_scheduler_(rpc_node_scheduler), op_context_(op_context) { - rpc_node_scheduler_->SetOpcontext(op_context_); + explicit RpcActorOpContextSetter(RpcNodeScheduler *rpc_node_scheduler, const RpcActorSetPtr &rpc_actors, + OpContext *const op_context) + : rpc_node_scheduler_(rpc_node_scheduler), rpc_actors_(rpc_actors), op_context_(op_context) { + rpc_node_scheduler_->SetOpcontext(rpc_actors_, op_context_); } - ~RpcActorOpContextSetter() { rpc_node_scheduler_->ResetOpcontext(); } + ~RpcActorOpContextSetter() { rpc_node_scheduler_->ResetOpcontext(rpc_actors_); } private: RpcNodeScheduler *rpc_node_scheduler_; + RpcActorSetPtr rpc_actors_; OpContext *op_context_; }; } // namespace runtime diff --git a/mindspore/python/mindspore/train/dataset_helper.py b/mindspore/python/mindspore/train/dataset_helper.py index 98623ef5220..7dc22fe6899 100644 --- a/mindspore/python/mindspore/train/dataset_helper.py +++ b/mindspore/python/mindspore/train/dataset_helper.py @@ -369,7 +369,7 @@ class _DatasetIter: if not hasattr(dataset, '__transfer_dataset__'): if hasattr(dataset, '__loop_size__'): # PS mode does not support loop sink and need get the real sink size. - if not _is_role_worker(): + if not (_is_role_worker() and _is_ps_mode()): self.sink_size = dataset.__loop_size__ create_data_info_queue = (sink_size == 1 and self.sink_count == 1 and dataset.get_dataset_size() != 1 and context.get_context("device_target") == "Ascend")