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

Replace DeviceIndexes with Devices in RRefs (#57442)

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

We did this for the RPC agents and for ivalue::Future, the last one (I think) is RRef.
ghstack-source-id: 128184664

Test Plan: CI

Reviewed By: mrshenli

Differential Revision: D28144368

fbshipit-source-id: eeacab6006f72118cbec542a02322f2e391c67a3
This commit is contained in:
Luca Wehrstedt
2021-05-06 01:10:10 -07:00
committed by Facebook GitHub Bot
parent 8e9bbd3113
commit 7ffadf6e46
8 changed files with 19 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@ -313,7 +313,7 @@ c10::intrusive_ptr<RRef> RRefContext::getOrCreateRRef(
c10::intrusive_ptr<OwnerRRef> RRefContext::getOrCreateOwnerRRef(
const RRefId& rrefId,
const TypePtr& type,
std::vector<c10::DeviceIndex> devices) {
std::vector<c10::Device> devices) {
std::lock_guard<std::mutex> lock(mutex_);
const auto iter = owners_.find(rrefId);
if (iter == owners_.end()) {
@ -369,7 +369,7 @@ c10::intrusive_ptr<OwnerRRef> RRefContext::getOrCreateOwnerRRef(
c10::intrusive_ptr<OwnerRRef> RRefContext::createOwnerRRef(
const TypePtr& type,
std::vector<c10::DeviceIndex> devices) {
std::vector<c10::Device> devices) {
// Don't add this OnwerRRef to the owners_ map yet, otherwise
// it will never be removed from there. Instead, only add it to the
// map in prepareChildFork, in case this local RRef is being passed
@ -395,6 +395,8 @@ c10::intrusive_ptr<JitFuture> RRefContext::getOwnerRRef(
// Note: The type passed into RRefType::create() does not matter here, as
// the future is marked as completed with the RRef of the specific type
// in getOrCreateOwnerRRef().
// Also no need to specify any devices because the RRef object itself
// doesn't contain any DataPtrs, it just provides means to retrieve them.
auto futureOwner =
c10::make_intrusive<JitFuture>(RRefType::create(c10::AnyType::get()));
pendingOwners_[rrefId] = futureOwner;
@ -408,6 +410,8 @@ c10::intrusive_ptr<JitFuture> RRefContext::getOwnerRRef(
auto owner = iter->second;
auto rrefPtr = fromOwnerRRef(owner);
// No need to specify any devices because the RRef object itself doesn't
// contain any DataPtrs, it just provides means to retrieve them.
auto futureOwner =
c10::make_intrusive<JitFuture>(RRefType::create(owner->type()));
futureOwner->markCompleted(IValue(rrefPtr));