Let RpcAgent::send() return JitFuture (#49906)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49906

This commit modifies RPC Message to inherit from `torch::CustomClassHolder`,
and wraps a Message in an IValue in `RpcAgent::send()`.

Test Plan: Imported from OSS

Reviewed By: lw

Differential Revision: D25719518

Pulled By: mrshenli

fbshipit-source-id: 694e40021e49e396da1620a2f81226522341550b
This commit is contained in:
Shen Li
2021-01-07 19:43:44 -08:00
committed by Facebook GitHub Bot
parent 4de6b279c8
commit 84e3237a53
13 changed files with 78 additions and 22 deletions

View File

@ -331,10 +331,10 @@ void PyRRef::backward(
// Invoke distributed backward remotely.
auto rpcAgent = rpc::RpcAgent::getCurrentRpcAgent();
rpcAgent
->send(
rpc::RpcAgent::toFutureMessage(
rpcAgent->send(
rpcAgent->getWorkerInfo(rref->owner()),
std::move(rrefBackwardReq).toMessage())
std::move(rrefBackwardReq).toMessage()))
->wait();
}
}