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