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
9namespace torch {
10namespace distributed {
11namespace rpc {
12
13class 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