mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
3fb1e73a4e
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/39758 Test Plan: Imported from OSS Differential Revision: D21963789 Pulled By: mrshenli fbshipit-source-id: f16f464ba01401b160cc4d3daf036e4bc806d7ea
42 lines
1.7 KiB
C++
42 lines
1.7 KiB
C++
#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
|