1 #pragma once
2 
3 #include <torch/csrc/distributed/autograd/rpc_messages/autograd_metadata.h>
4 #include <torch/csrc/distributed/rpc/message.h>
5 #include <torch/csrc/distributed/rpc/rpc_command_base.h>
6 
7 namespace torch {
8 namespace distributed {
9 namespace autograd {
10 
11 // Used to request other workers to clean up their autograd context.
12 class TORCH_API CleanupAutogradContextReq : public rpc::RpcCommandBase {
13  public:
14   explicit CleanupAutogradContextReq(int64_t context_id);
15   // Serialization and deserialization methods.
16   c10::intrusive_ptr<rpc::Message> toMessageImpl() && override;
17   static std::unique_ptr<CleanupAutogradContextReq> fromMessage(
18       const rpc::Message& message);
19 
20   // Retrieve the context id we are cleaning up with this message.
21   int64_t getContextId();
22 
23  private:
24   int64_t context_id_;
25 };
26 
27 } // namespace autograd
28 } // namespace distributed
29 } // namespace torch
30