frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[RPC] Support timeout in rref._get_type() (#50498)

Browse Source 
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50498

This change is mostly needed for the next diff in this stack, where
rref._get_type() is called in the rpc_async/rpc_sync RRef proxy function and
can block indefinitely if there is no timeout. It will also be useful to have a
timeout argument when we publicize this API to keep it consistent with other
RPC APIs.
ghstack-source-id: 119859767

Test Plan: Added UT

Reviewed By: pritamdamania87

Differential Revision: D25897588

fbshipit-source-id: 2e84aaf7e4faecf80005c78ee2ac8710f387503e
This commit is contained in:
Rohan Varma
2021-01-15 13:16:15 -08:00
committed by Facebook GitHub Bot
parent c78e7db7ee
commit ab1ba8f433
5 changed files with 35 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
torch/csrc/distributed/rpc/py_rref.cpp
View File

@ -249,14 +249,15 @@ py::object PyRRef::createRRefProxy(const RRefProxyType& type) const {
}
}
py::object PyRRef::getRRefType() {
py::object PyRRef::getRRefType(float timeout) {
// GIL is not released when calling this function.
if (!type_.has_value()) {
pybind11::gil_scoped_release release;
auto& pythonRpcHandler = PythonRpcHandler::getInstance();
auto& typeFuncs = pythonRpcHandler.getRRefTypeFunctions();
pybind11::gil_scoped_acquire acquire;
type_ = isOwner() ? typeFuncs.onOwner_(*this) : typeFuncs.onUser_(*this);
type_ = isOwner() ? typeFuncs.onOwner_(*this)
: typeFuncs.onUser_(*this, timeout);
}
return *type_;