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 { |
10 | namespace distributed { |
11 | namespace rpc { |
12 | |
13 | class TORCH_API PythonRemoteCall : public RpcCommandBase { |
14 | public: |
15 | PythonRemoteCall( |
16 | SerializedPyObj&& serializedPyObj, |
17 | at::IValue retRRefId, |
18 | at::IValue retForkId, |
19 | const bool isAsyncExecution); |
20 | |
21 | inline const SerializedPyObj& serializedPyObj() const { |
22 | return serializedPyObj_; |
23 | } |
24 | |
25 | inline const at::IValue& retRRefId() const { |
26 | return retRRefId_; |
27 | } |
28 | |
29 | inline const at::IValue& retForkId() const { |
30 | return retForkId_; |
31 | } |
32 | |
33 | inline bool isAsyncExecution() const { |
34 | return isAsyncExecution_; |
35 | } |
36 | |
37 | c10::intrusive_ptr<Message> toMessageImpl() && override; |
38 | static std::unique_ptr<PythonRemoteCall> fromMessage(const Message& message); |
39 | |
40 | private: |
41 | SerializedPyObj serializedPyObj_; |
42 | const at::IValue retRRefId_; |
43 | const at::IValue retForkId_; |
44 | const bool isAsyncExecution_; |
45 | }; |
46 | |
47 | } // namespace rpc |
48 | } // namespace distributed |
49 | } // namespace torch |
50 |