Acquire GIL when constructing/destructing ConcretePyObjectHolder (#37870)

mrshenli · facebook-github-bot · commit ee1ddcef8d29 · 2020-05-07T07:37:39.000-07:00
Summary: Pull Request resolved: #37870 Test Plan: Imported from OSS Differential Revision: D21410785 fbshipit-source-id: 374d5f40fbdfec98262aa4c84ec4ccdc40fb2ac1
diff --git a/torch/csrc/distributed/rpc/python_functions.cpp b/torch/csrc/distributed/rpc/python_functions.cpp
@@ -105,52 +105,29 @@ std::shared_ptr<FutureMessage> sendPythonRemoteCall(
       true /*forceGradRecording*/);
 }
 
-void DeleteFutureIValue(FutureIValue* fv) {
-  if (fv->constValue().isPyObject()) {
-    pybind11::gil_scoped_acquire ag;
-    delete fv;
-  } else {
-    delete fv;
-  }
-}
-
 } // namespace
 
 using namespace torch::distributed::autograd;
 
 std::shared_ptr<FutureIValue> toFutureIValue(
     const std::shared_ptr<FutureMessage>& fm,
     bool hasValue) {
-  if (hasValue) {
-    // NB: The custom deleter is necessary because the FutureIValue object
-    // holds a py::object and it would require GIL to delete.
-    std::shared_ptr<FutureIValue> fv(new FutureIValue(), DeleteFutureIValue);
-
-    fm->addCallback([fv](const FutureMessage& fm) {
-      // Don't need to acquire GIL here, as toPyObj acquires GIL
-      // when creating the py::object
-      if (fm.hasError()) {
-        fv->setError(*fm.error());
-      } else {
-        fv->markCompleted(
-            jit::toIValue(toPyObj(fm.constValue()), PyObjectType::get()));
-      }
-    });
-
-    return fv;
-  } else {
-    auto fv = std::make_shared<FutureIValue>();
-
-    fm->addCallback([fv](const FutureMessage& fm) {
-      if (fm.hasError()) {
-        fv->setError(*fm.error());
-      } else {
-        fv->markCompleted(IValue());
-      }
-    });
+  auto fv = std::make_shared<FutureIValue>();
+
+  fm->addCallback([fv, hasValue](const FutureMessage& fm) {
+    if (fm.hasError()) {
+      fv->setError(*fm.error());
+    } else if (hasValue) {
+      // Don't need to acquire GIL here, as toPyObj and toIValue acquires GIL
+      // when creating/copying the py::object
+      fv->markCompleted(
+          jit::toIValue(toPyObj(fm.constValue()), PyObjectType::get()));
+    } else {
+      fv->markCompleted(IValue());
+    }
+  });
 
-    return fv;
-  }
+  return fv;
 }
 
 std::shared_ptr<FutureIValue> pyRpcBuiltin(
diff --git a/torch/csrc/jit/python/pybind_utils.h b/torch/csrc/jit/python/pybind_utils.h
@@ -604,11 +604,9 @@ inline IValue toIValue(
       AT_ERROR("RRef is only supported with the distributed package");
 #endif
     } break;
-    case TypeKind::PyObjectType:
-      // convert a py::handle to the IValue that holds the py::object
-      return c10::ivalue::ConcretePyObjectHolder::create(
-          obj.cast<py::object>());
-
+    case TypeKind::PyObjectType: {
+      return c10::ivalue::ConcretePyObjectHolder::create(obj);
+    }
     case TypeKind::CapsuleType: {
       return IValue::make_capsule(
           py::cast<c10::intrusive_ptr<CustomClassHolder>>(obj));
diff --git a/torch/csrc/jit/python/python_ivalue.h b/torch/csrc/jit/python/python_ivalue.h
@@ -11,17 +11,27 @@ namespace ivalue {
 struct C10_EXPORT ConcretePyObjectHolder final : PyObjectHolder {
  public:
   static c10::intrusive_ptr<PyObjectHolder> create(py::object py_obj) {
-    return c10::make_intrusive<ConcretePyObjectHolder>(py_obj);
+    return c10::make_intrusive<ConcretePyObjectHolder>(std::move(py_obj));
+  }
+
+  static c10::intrusive_ptr<PyObjectHolder> create(const py::handle& handle) {
+    py::gil_scoped_acquire ag;
+    return c10::make_intrusive<ConcretePyObjectHolder>(
+        handle.cast<py::object>());
   }
 
   PyObject* getPyObject() override {
     return py_obj_.ptr();
   }
 
-  ~ConcretePyObjectHolder() {}
+  ~ConcretePyObjectHolder() {
+    pybind11::gil_scoped_acquire ag;
+    py_obj_ = py::none();
+  }
   // explicit construction to avoid errornous implicit conversion and
   // copy-initialization
-  explicit ConcretePyObjectHolder(py::object py_obj) : py_obj_(py_obj) {}
+  explicit ConcretePyObjectHolder(py::object py_obj)
+      : py_obj_(std::move(py_obj)) {}
 
  private:
   py::object py_obj_;
