xref: /aosp_15_r20/external/pytorch/torch/csrc/distributed/rpc/unpickled_python_call.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <torch/csrc/distributed/rpc/unpickled_python_call.h>
2 
3 #include <torch/csrc/distributed/rpc/python_rpc_handler.h>
4 
5 namespace torch::distributed::rpc {
6 
UnpickledPythonCall(const SerializedPyObj & serializedPyObj,bool isAsyncExecution)7 UnpickledPythonCall::UnpickledPythonCall(
8     const SerializedPyObj& serializedPyObj,
9     bool isAsyncExecution)
10     : isAsyncExecution_(isAsyncExecution) {
11   auto& pythonRpcHandler = PythonRpcHandler::getInstance();
12   pybind11::gil_scoped_acquire ag;
13   pythonUdf_ = pythonRpcHandler.deserialize(serializedPyObj);
14 }
15 
~UnpickledPythonCall()16 UnpickledPythonCall::~UnpickledPythonCall() {
17   // explicitly setting PyObject* to nullptr to prevent py::object's dtor to
18   // decref on the PyObject again.
19   // See Note [Destructing py::object] in python_ivalue.h
20   py::gil_scoped_acquire acquire;
21   pythonUdf_.dec_ref();
22   pythonUdf_.ptr() = nullptr;
23 }
24 
toMessageImpl()25 c10::intrusive_ptr<Message> UnpickledPythonCall::toMessageImpl() && {
26   TORCH_INTERNAL_ASSERT(
27       false, "UnpickledPythonCall does not support toMessage().");
28 }
29 
pythonUdf() const30 const py::object& UnpickledPythonCall::pythonUdf() const {
31   return pythonUdf_;
32 }
33 
34 } // namespace torch::distributed::rpc
35