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/distributed/rpc/unpickled_python_call.h> 6 #include <torch/csrc/utils/pybind.h> 7 8 namespace torch::distributed::rpc { 9 10 // This class converts the content in a PythonRemoteCall into py::object. This 11 // is a helper class to make sure that all arguments deserialization is done 12 // before entering RequestCallbackImpl::processRpc(...), so that the 13 // deserialization related logic can be carried out in one spot instead of 14 // scattered in multiple places for different message types. 15 // NB: The reason for not consolidating class into PythonRemoteCall is because 16 // PythonRemoteCall is a libtorch type which should not depend on Python types. 17 class TORCH_API UnpickledPythonRemoteCall final : public UnpickledPythonCall { 18 public: 19 explicit UnpickledPythonRemoteCall( 20 const SerializedPyObj& serializedPyObj, 21 const at::IValue& retRRefId, 22 const at::IValue& retForkId, 23 const bool isAsyncExecution); 24 25 const RRefId& rrefId() const; 26 const ForkId& forkId() const; 27 28 private: 29 RRefId rrefId_; 30 ForkId forkId_; 31 }; 32 33 } // namespace torch::distributed::rpc 34