xref: /aosp_15_r20/external/pytorch/torch/csrc/distributed/rpc/unpickled_python_call.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
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