[RPC] Add to confirmed users immediately if the fork is shared from owner, instead of adding nothing to pending users (#34988)

xush6528 · facebook-github-bot · commit e5ee95e448f4 · 2020-03-18T18:17:41.000-07:00
Summary: Pull Request resolved: #34988 In #31893, we introduced a confirmedUsers_ map in RRefContext. For the case that the fork is shared from the owner, there is no `pendingUsers_` intermediate phase for this fork, we should put this fork into `confirmedUsers_` immediately. Test Plan: ``` buck test mode/dev-nosan //caffe2/test/distributed/rpc:rpc_fork ``` ``` buck test mode/dev-nosan //caffe2/test/distributed/rpc/jit:rpc_fork ``` Differential Revision: D7735909 fbshipit-source-id: 14c36a16486f0cc9618dcfb111fe5223781b647d
diff --git a/torch/csrc/distributed/rpc/rref_context.cpp b/torch/csrc/distributed/rpc/rref_context.cpp
@@ -344,6 +344,7 @@ void RRefContext::notifyOwnerAndParentOfFork(
     const ForkId& forkId,
     worker_id_t parent,
     const c10::intrusive_ptr<RRef>& rref) {
+  // Fork is shared from owner.
   if (parent == rref->owner()) {
     if (parent == agent_->getWorkerInfo().id_) {
       // Owner sending RRef to self, remove the forkId as it was added during
@@ -365,10 +366,13 @@ void RRefContext::notifyOwnerAndParentOfFork(
       // Hence, it is not necessary to send another RREF_CHILD_ACCEPT or
       // RREF_FORK_REQUEST back to the owner. See Note [Early Fork
       // Registration].
+      std::lock_guard<std::mutex> lock(mutex_);
+      addConfirmedUser(forkId, rref);
     }
     return;
   }
 
+  // Fork is shared from user.
   if (rref->isOwner()) {
     // See Note [Useful Phantom Fork ID for User to Owner Call]
     // In this case, the owner is the caller, and it does not add the fork id
@@ -494,17 +498,25 @@ void RRefContext::delPendingUser(const ForkId& forkId) {
     // hiding the subtle logic using a reentrant lock.
     deletedState = iter->second; // Increase refcount
 
-    confirmedUsers_.emplace(
-        std::piecewise_construct,
-        std::forward_as_tuple(forkId),
-        std::forward_as_tuple(iter->second->rref_));
+    addConfirmedUser(forkId, iter->second->rref_);
     pendingUsers_.erase(iter); // Decrease refcount.
   }
   deletedState->confirm();
   deleteAllUsersCV_.notify_all();
   deletedState.reset(); // Decrease refcount.
 }
 
+void RRefContext::addConfirmedUser(
+    const ForkId& forkId,
+    const c10::intrusive_ptr<RRef>& rref) {
+  // Notice, caller need to hold the mutex for confirmedUsers_.
+  // std::lock_guard<std::mutex> lock(mutex_);
+  confirmedUsers_.emplace(
+      std::piecewise_construct,
+      std::forward_as_tuple(forkId),
+      std::forward_as_tuple(rref));
+}
+
 void RRefContext::recordThreadLocalPendingRRefs() {
   TORCH_INTERNAL_ASSERT(
       userTable_.empty(),
diff --git a/torch/csrc/distributed/rpc/rref_context.h b/torch/csrc/distributed/rpc/rref_context.h
@@ -147,6 +147,9 @@ class TORCH_API RRefContext {
       const ForkId& forkId,
       const c10::intrusive_ptr<RRef>& rref);
   void delPendingUser(const ForkId& forkId);
+  void addConfirmedUser(
+      const ForkId& forkId,
+      const c10::intrusive_ptr<RRef>& rref);
 
   // Start recroding new pending UserRRefs. All pending UserRRefs introduced
   // after this point will be put into the thread_local userTable_, which will
