!31937 Fix eval after train issue.

Merge pull request !31937 from ZPaC/add-dist-execution-mode

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<SplitGraphSegment> 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);
 }
 

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<RpcNodeScheduler>();
   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<DeviceCont
 
 #ifdef ENABLE_RPC_ACTOR
   // Set OpContext to rpc node scheduler.
-  auto op_context_setter = std::make_shared<RpcActorOpContextSetter>(rpc_node_scheduler_.get(), &op_context);
+  auto op_context_setter =
+    std::make_shared<RpcActorOpContextSetter>(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<SendActor>(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<RecvActor>(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.";

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<RpcActorSet>();
-  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<KernelActorPtr> kernel_actors = actor_set->kernel_actors_;
+  RpcActorSetPtr rpc_actor_set = std::make_shared<RpcActorSet>();
+  MS_EXCEPTION_IF_NULL(rpc_actor_set);
 
-RpcActorSetPtr RpcNodeScheduler::Build(const GraphCompilerInfo &) {
-  MS_EXCEPTION_IF_NULL(rpc_actor_set_);
   std::vector<RpcActorPtr> 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<RpcActor>(kernel_actor);
+    if (std::dynamic_pointer_cast<RpcActor>(kernel_actor) == nullptr) {
+      continue;
+    } else {
+      rpc_actors.emplace_back(rpc_actor);
+      if (std::dynamic_pointer_cast<SendActor>(rpc_actor) != nullptr) {
+        rpc_actor_set->send_actors_.emplace_back(std::dynamic_pointer_cast<SendActor>(rpc_actor));
+      } else if (std::dynamic_pointer_cast<RecvActor>(rpc_actor) != nullptr) {
+        rpc_actor_set->recv_actors_.emplace_back(std::dynamic_pointer_cast<RecvActor>(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<SendActorPtr> send_actors = rpc_actor_set_->send_actors_;
-  std::vector<RecvActorPtr> 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<SendActorPtr> send_actors = rpc_actor_set->send_actors_;
+  std::vector<RecvActorPtr> 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<DeviceTensor> *const op_context) {
+void RpcNodeScheduler::SetOpcontext(const RpcActorSetPtr &rpc_actors, OpContext<DeviceTensor> *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();
   }

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<DeviceTensor> *const op_context);
+  void SetOpcontext(const RpcActorSetPtr &rpc_actors, OpContext<DeviceTensor> *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<DeviceTensor> *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<DeviceTensor> *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<DeviceTensor> *op_context_;
 };
 }  // namespace runtime

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")