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] Add option to make rref.get_type not block. (#50977)

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

Adds a `blocking` flag that can be set to False to make this API return a `Future` to the type. This is to make this function non-blocking, mostly for a future change that will allow `rref.rpc_async()` to be completely non-blocking (it currently calls and waits for this function that issues an RPC in-line).
ghstack-source-id: 121021433

Test Plan: Modified UT

Reviewed By: mrshenli

Differential Revision: D25944582

fbshipit-source-id: e3b48a52af2d4578551a30ba6838927b489b1c03
This commit is contained in:
Rohan Varma
2021-02-04 20:15:25 -08:00
committed by Facebook GitHub Bot
parent 716a8c2153
commit c3f2f3294e
5 changed files with 123 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -251,17 +251,17 @@ py::object PyRRef::createRRefProxy(
}
}
py::object PyRRef::getRRefType(float timeout) {
py::object PyRRef::getRRefType(float timeout, bool blocking) {
// 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, timeout);
type_ = isOwner() ? typeFuncs.onOwner_(*this, blocking)
: typeFuncs.onUser_(*this, timeout, blocking);
}
// Returns py::object that can be Python type or future.
return *type_;
}