fix python rpc handler exit crash (pytorch#27251)

Summary:
Pull Request resolved: pytorch#27251

 Explicitly clean up py::objects to avoid segment faults when py::objects with CPython are cleaned up later at program exit.

See similar issues reported pybind/pybind11#1598
and pybind/pybind11#1493.

Our local tests also caught this segment faults if py::objects are cleaned
up at program exit. The explaination is: CPython cleans up most critical
utitlies before cleaning up PythonRpcHandler singleton, so when
PythonRpcHandler signleton cleans up py::objects and call dec_ref(), it
will crash.

The solution is to clean up py::objects earlier when Rpc agent join().
Be note that py::objects can not be cleaned up when Rpc agent is destroyed
as well, as Rpc agent is global variable and it will have same issue as
PythonRpcHandler.

close pytorch#27182
ghstack-source-id: 92035069

Test Plan: unit tests on python 3.6 and python 3.5

Differential Revision: D17727362

fbshipit-source-id: c254023f6a85acce35528ba756a4efabba9a519f

Author: zhaojuanmao

test/run_test.py

@@ -16,6 +16,7 @@
 from torch.utils import cpp_extension
 from common_utils import TEST_WITH_ROCM, shell
 import torch.distributed as dist
+PY33 = sys.version_info >= (3, 3)
 PY36 = sys.version_info >= (3, 6)
 
 TESTS = [
@@ -61,22 +62,22 @@
     'function_schema',
 ]
 
-# skip < 3.6 b/c fstrings added in 3.6 for jit_py3
-# skip < 3.6 for rpc_spawn and dist_autograd_spawn temporarily because
-# a segmenation fault was triggered on python 3.5,
-# rpc_spawn and dist_autograd_spawn tests were added in
-# https://github.com/pytorch/pytorch/pull/25656
-# skip < 3.6 for rpc_fork as it imports mock that is only available in 3.6, mock
-# was added to rpc_fork in https://github.com/pytorch/pytorch/pull/26997
-if PY36:
+# skip < 3.3 because mock is added in 3.3 and is used in rpc_fork and rpc_spawn
+# skip python2 for rpc and dist_autograd tests that do not support python2
+if PY33:
     TESTS.extend([
-        'jit_py3',
         'rpc_fork',
         'rpc_spawn',
         'dist_autograd_fork',
         'dist_autograd_spawn',
     ])
 
+# skip < 3.6 b/c fstrings added in 3.6
+if PY36:
+    TESTS.extend([
+        'jit_py3',
+    ])
+
 WINDOWS_BLACKLIST = [
     'distributed',
     'rpc_fork',

torch/csrc/distributed/rpc/init.cpp

@@ -71,6 +71,9 @@ PyObject* rpc_init(PyObject* /* unused */) {
                 return PyRRef::unpickle(t);
               }));
 
+  // future.wait() should not be called after join_rpc(), e.g., pythonRpcHandler
+  // is cleaned up in join_rpc(), after join_rpc(), python objects returned
+  // from rpc python call can not be resolved.
   auto futureMessage =
       shared_ptr_class_<FutureMessage>(module, "FutureMessage")
           .def(
@@ -112,6 +115,10 @@ PyObject* rpc_init(PyObject* /* unused */) {
     RRefContext::getInstance().destroyInstance();
   });
 
+  module.def("_cleanup_python_rpc_handler", []() {
+    PythonRpcHandler::getInstance().cleanup();
+  });
+
   module.def(
       "invoke_rpc_builtin",
       [](RpcAgent& agent,

torch/csrc/distributed/rpc/process_group_agent.cpp

@@ -147,9 +147,6 @@ ProcessGroupAgent::ProcessGroupAgent(
   for (int rank = 0; rank < (int)tmpWorkerIds.size(); ++rank) {
     allWorkerInfo_.emplace_back(std::move(tmpWorkerIds[rank]), rank);
   }
-
-  // construct PythonRpcHandler singleton here
-  PythonRpcHandler::getInstance();
 }
 
 const WorkerInfo& ProcessGroupAgent::getWorkerInfo(

torch/csrc/distributed/rpc/python_rpc_handler.cpp

@@ -26,6 +26,13 @@ PythonRpcHandler::PythonRpcHandler() {
   pySerialize_ = getFunction(module, "serialize");
 }
 
+void PythonRpcHandler::cleanup() {
+  AutoGIL ag;
+  pyRunFunction_ = py::none();
+  pyLoadReturnValue_ = py::none();
+  pySerialize_ = py::none();
+}
+
 PythonRpcHandler& PythonRpcHandler::getInstance() {
   static PythonRpcHandler handler;
   return handler;

torch/csrc/distributed/rpc/python_rpc_handler.h

@@ -39,6 +39,21 @@ class PYBIND11_EXPORT PythonRpcHandler {
   // Deserialize a string into a py::object
   py::object deserialize(const SerializedPyObj& serializedObj);
 
+  // Explicitly clean up py::objects to avoid segment faults when
+  // py::objects with CPython are cleaned up later at program exit
+  // See similar issues reported https://github.com/pybind/pybind11/issues/1598
+  // and https://github.com/pybind/pybind11/issues/1493
+  // Our local tests also caught this segment faults if py::objects are cleaned
+  // up at program exit. The explaination is: CPython cleans up most critical
+  // utilities before cleaning up PythonRpcHandler singleton, so when
+  // PythonRpcHandler signleton cleans up py::objects and call dec_ref(), it
+  // will crash.
+  // The solution is to clean up py::objects earlier when Rpc agent join().
+  // Be note that py::objects can not be cleaned up when Rpc agent is destroyed
+  // as well, as Rpc agent is global variable and it will have same issue as
+  // PythonRpcHandler.
+  void cleanup();
+
  private:
   PythonRpcHandler();
   ~PythonRpcHandler() = default;

torch/distributed/rpc/api.py

@@ -1,7 +1,7 @@
 from torch.distributed import invoke_rpc_builtin, invoke_rpc_python_udf
 from torch.distributed import invoke_remote_builtin, invoke_remote_python_udf
 from torch.distributed import _start_rpc_agent
-from torch.distributed import _destroy_rref_context
+from torch.distributed import _destroy_rref_context, _cleanup_python_rpc_handler
 from torch.distributed import ProcessGroupAgent
 from torch.distributed import WorkerInfo
 from .backend_registry import is_backend_registered, init_backend
@@ -40,6 +40,11 @@ def join_rpc():
         _agent.join()
         _agent = None
         _destroy_rref_context()
+        # clean up python rpc handler in join_rpc(), see comments in
+        # PythonRpcHandler::cleanup(), call it in python API because the
+        # cleanup() function has python dependency, it assumes python
+        # interpreter exists
+        _cleanup_python_rpc_handler()
 
 
 @_require_initialized