xref: /aosp_15_r20/external/pytorch/torch/csrc/distributed/rpc/python_remote_call.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/csrc/distributed/rpc/message.h>
4 #include <torch/csrc/distributed/rpc/rpc_command_base.h>
5 #include <torch/csrc/distributed/rpc/types.h>
6 #include <torch/csrc/jit/serialization/pickler.h>
7 #include <vector>
8 
9 namespace torch::distributed::rpc {
10 
11 class TORCH_API PythonRemoteCall : public RpcCommandBase {
12  public:
13   PythonRemoteCall(
14       SerializedPyObj&& serializedPyObj,
15       at::IValue retRRefId,
16       at::IValue retForkId,
17       const bool isAsyncExecution);
18 
serializedPyObj()19   inline const SerializedPyObj& serializedPyObj() const {
20     return serializedPyObj_;
21   }
22 
retRRefId()23   inline const at::IValue& retRRefId() const {
24     return retRRefId_;
25   }
26 
retForkId()27   inline const at::IValue& retForkId() const {
28     return retForkId_;
29   }
30 
isAsyncExecution()31   inline bool isAsyncExecution() const {
32     return isAsyncExecution_;
33   }
34 
35   c10::intrusive_ptr<Message> toMessageImpl() && override;
36   static std::unique_ptr<PythonRemoteCall> fromMessage(const Message& message);
37 
38  private:
39   SerializedPyObj serializedPyObj_;
40   const at::IValue retRRefId_;
41   const at::IValue retForkId_;
42   const bool isAsyncExecution_;
43 };
44 
45 } // namespace torch::distributed::rpc
46