1 #pragma once 2 3 #include <ATen/core/ivalue.h> 4 #include <torch/csrc/autograd/profiler.h> 5 #include <torch/csrc/distributed/autograd/utils.h> 6 #include <torch/csrc/distributed/rpc/rref_context.h> 7 #include <torch/csrc/distributed/rpc/script_remote_call.h> 8 9 namespace torch::distributed::rpc { 10 11 // This function sends an rpc call to run torchscript function, currently the 12 // torchscript function could only be a user defined python function with 13 // "@torch.jit.script" annotation. The torchscript function could not be 14 // a class constructor, class method, instance method or a script module. 15 // dst: destination worker name 16 // qualifiedName: torchscript function qualified name string like 17 // "moduleName::torchscriptFunctionName", e.g, 18 // "dist_autograd_test::my_py_add" 19 // stack: a bag of IValue args passed to torchscriptFunctionName 20 // It returns c10::intrusive_ptr<ivalue::Future> 21 c10::intrusive_ptr<c10::ivalue::Future> TORCH_API rpcTorchscript( 22 const std::string& dstWorkerName, 23 const c10::QualifiedName& qualifiedName, 24 const c10::FunctionSchema& functionSchema, 25 std::vector<c10::IValue>& stack, 26 const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout, 27 const bool isAsyncExecution = false); 28 29 c10::intrusive_ptr<RRef> TORCH_API remoteTorchscript( 30 const std::string& dstWorkerName, 31 const c10::QualifiedName& qualifiedName, 32 const c10::FunctionSchema& functionSchema, 33 std::vector<c10::IValue>& stack, 34 const float rpcTimeoutSeconds = torch::distributed::rpc::kUnsetRpcTimeout, 35 const bool isAsyncExecution = false); 36 37 } // namespace torch::distributed::rpc 38