1 #pragma once 2 3 #include <torch/csrc/distributed/rpc/rpc_command_base.h> 4 #include <torch/csrc/distributed/rpc/types.h> 5 #include <torch/csrc/utils/pybind.h> 6 7 namespace torch::distributed::rpc { 8 9 // This class converts the content in a PythonCall into py::object. This is a 10 // helper class to make sure that all arguments deserialization is done before 11 // entering RequestCallbackImpl::processRpc(...), so that the deserialization 12 // related logic can be carried out in one spot instead of scattered in multiple 13 // places for different message types. 14 // NB: The reason for not consolidating class into PythonCall is because 15 // PythonCall is a libtorch type which should not depend on Python types. 16 class TORCH_API UnpickledPythonCall : public RpcCommandBase { 17 public: 18 UnpickledPythonCall( 19 const SerializedPyObj& serializedPyObj, 20 bool isAsyncExecution); 21 ~UnpickledPythonCall() override; 22 23 // toMessage() method is not implemented, as objects of this class should 24 // never be directly converted into a Message object. 25 c10::intrusive_ptr<Message> toMessageImpl() && override; 26 const py::object& pythonUdf() const; 27 isAsyncExecution()28 inline bool isAsyncExecution() const { 29 return isAsyncExecution_; 30 } 31 32 private: 33 py::object pythonUdf_; 34 // NOLINTNEXTLINE(cppcoreguidelines-avoid-const-or-ref-data-members) 35 const bool isAsyncExecution_; 36 }; 37 38 } // namespace torch::distributed::rpc 39