forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
rref_context.cpp
420 lines (384 loc) · 14.9 KB
/
rref_context.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
#include <torch/csrc/distributed/rpc/rref_context.h>
#include <torch/csrc/distributed/rpc/rref_proto.h>
#include <sstream>
namespace torch {
namespace distributed {
namespace rpc {
namespace callback {
void confirmPendingUser(
const rpc::Message& message,
const c10::optional<utils::FutureError>& futErr) {
RRefContext::handleException(futErr);
auto rr = RemoteRet::fromMessage(message);
auto& ctx = RRefContext::getInstance();
ctx.delPendingUser(rr->forkId());
}
} // namespace callback
// Keys for RRef-related debug information.
const std::string kNumOwnerRRefs = "num_owner_rrefs";
const std::string kNumPendingUsers = "num_pending_users";
RRefContext& RRefContext::getInstance() {
// Leaky singleton to avoid module destructor races.
static RRefContext* context = new RRefContext(RpcAgent::getCurrentRpcAgent());
return *context;
}
std::vector<std::shared_ptr<RRef>> RRefContext::destroyInstance(
bool ignoreRRefLeak) {
auto& ctx = RRefContext::getInstance();
{
std::lock_guard<std::mutex> lock(ctx.destroyedMutex_);
ctx.destroyed_ = true;
}
ctx.checkRRefLeaks(ignoreRRefLeak);
std::vector<std::shared_ptr<RRef>> deletedRRefs;
for (auto& entry : ctx.owners_) {
auto rref = entry.second;
if (rref->isPyObj()) {
deletedRRefs.emplace_back(std::move(rref));
}
}
ctx.owners_.clear();
return deletedRRefs;
}
void RRefContext::handleException(
const c10::optional<utils::FutureError>& futErr) {
if (futErr) {
// TODO: allow users to register an error handler and call it here.
VLOG(1) << "Got exception: " << (*futErr).what();
throw std::runtime_error((*futErr).what());
}
}
RRefContext::RRefContext(std::shared_ptr<RpcAgent> agent)
: agent_(std::move(agent)), destroyed_(false) {}
RRefContext::~RRefContext() {
if (!owners_.empty()) {
VLOG(1) << "Destructing RRefContext with non-empty OwnerRRef set. "
<< "This would likely cause Python deref error. "
<< "Make sure destroyInstance() is invoked before destruction.";
}
}
std::unordered_map<std::string, std::string> RRefContext::getDebugInfo() {
std::unordered_map<std::string, std::string> info;
std::unique_lock<std::mutex> lock(mutex_);
auto ownerSize = owners_.size();
auto numPendingUsers = pendingUsers_.size();
lock.unlock();
info[kNumOwnerRRefs] = c10::to_string(ownerSize);
info[kNumPendingUsers] = c10::to_string(numPendingUsers);
return info;
}
void RRefContext::checkRRefLeaks(bool ignoreRRefLeak) {
if (!forks_.empty()) {
std::stringstream ss;
for (auto& entry : forks_) {
const RRefId& rrefId = entry.first;
for (const auto& forkId : entry.second) {
ss << "Leaking RRef " << rrefId << " with fork Id " << forkId
<< std::endl;
}
}
LOG(WARNING)
<< "Detected RRef Leaks during shutdown. This usually "
<< "occurs when the application code still holds references to RRef "
<< "instances when calling shutdown(). If the program has "
<< "completed correctly and the process is exiting, it is OK to "
<< "ignore these leaks. However, if you program will keep running "
<< "after this, these leaks could result in memory leaks on RRef "
<< "owners. Please make sure all RRefs are out of scope and Python "
<< "GC has deleted them before calling shutdown(): \n"
<< ss.str();
if (!ignoreRRefLeak) {
TORCH_CHECK(false, ss.str());
}
}
}
std::shared_ptr<UserRRef> RRefContext::createUserRRef(
worker_id_t ownerId,
const TypePtr& type) {
TORCH_CHECK(ownerId != getWorkerId(), "Cannot create UserRRef on owner.");
// Explicitly creating rrefId before forkId to make sure the order is
// deterministic, as the argument evaluation order is system and compiler
// dependent.
const auto rrefId = genGloballyUniqueId();
const auto forkId = genGloballyUniqueId();
return createUserRRef(ownerId, rrefId, forkId, type);
}
std::shared_ptr<UserRRef> RRefContext::createUserRRef(
worker_id_t ownerId,
const RRefId& rrefId,
const ForkId& forkId,
const TypePtr& type) {
TORCH_CHECK(ownerId != getWorkerId(), "RRef owner cannot create user RRef.");
// RRefContext does not track user RRefs, it will be destructed when there
// is no shared_ptrs pointing to it.
//
// NB: cannot use make_shared here as the constructor of UserRRef is private.
// NB: This UserRRef has not been confirmed by the owner yet. This function's
// call site is responsible for adding this UserRRef to pendingUsers_.
// Currently, there are two call sites.
// (1) The creator user in python_functions.cpp
// (2) The callee user in RRefContext::notifyOwnerAndParentOfFork.
//
// The reason for not adding the pending user here is to put addPendingUser()
// close to where the RPC occurs, and it is more clear to pair it with
// deletePendingUser() in the response callback at the call site.
return std::shared_ptr<UserRRef>(new UserRRef(ownerId, rrefId, forkId, type));
}
void RRefContext::delUser(
const worker_id_t owner,
const RRefId& rrefId,
const ForkId& forkId) {
std::lock_guard<std::mutex> lock(destroyedMutex_);
if (!destroyed_) {
auto fm = agent_->send(
agent_->getWorkerInfo(owner),
RRefUserDelete(rrefId, forkId).toMessage());
fm->addCallback([](const Message& /* unused */,
const c10::optional<utils::FutureError>& futErr) {
RRefContext::handleException(futErr);
});
}
}
std::shared_ptr<RRef> RRefContext::getOrCreateRRef(
const RRefForkData& rfd,
const TypePtr& type) {
auto& ownerId = rfd.ownerId_;
auto& rrefId = rfd.rrefId_;
auto& forkId = rfd.forkId_;
if (ownerId == getWorkerId()) {
auto ownerRRef = getOwnerRRef(rrefId);
TORCH_INTERNAL_ASSERT(ownerRRef->type() == type);
return ownerRRef;
} else {
return createUserRRef(ownerId, rrefId, forkId, type);
}
}
std::shared_ptr<OwnerRRef> RRefContext::getOrCreateOwnerRRef(
const RRefId& rrefId,
const TypePtr& type) {
std::lock_guard<std::mutex> lock(mutex_);
const auto iter = owners_.find(rrefId);
if (iter == owners_.end()) {
// Scenario (1) the first time this owner knows about this RRef
//
// NB: cannot use make_shared here as the constructor of OwnerRRef is
// private.
auto rref =
std::shared_ptr<OwnerRRef>(new OwnerRRef(getWorkerId(), rrefId, type));
owners_[rref->rrefId()] = rref;
ownerCV_.notify_all();
return rref;
} else {
// Scenario (2) retrieving an existing RRef
auto ownerRRef = std::static_pointer_cast<OwnerRRef>(iter->second);
TORCH_INTERNAL_ASSERT(ownerRRef->type() == type);
return ownerRRef;
}
}
std::shared_ptr<OwnerRRef> RRefContext::createOwnerRRef(const TypePtr& type) {
// 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
// to another worker.
return std::shared_ptr<OwnerRRef>(
new OwnerRRef(getWorkerId(), genGloballyUniqueId(), type));
}
std::shared_ptr<OwnerRRef> RRefContext::getOwnerRRef(const RRefId& rrefId) {
std::unique_lock<std::mutex> lock(mutex_);
const auto iter = owners_.find(rrefId);
if (iter == owners_.end()) {
// Scenario (1) RRef is used before it is created
ownerCV_.wait(lock, [&] { return owners_.find(rrefId) != owners_.end(); });
return std::static_pointer_cast<OwnerRRef>(owners_[rrefId]);
} else {
// Scenario (2) retrieving an existing RRef
return std::static_pointer_cast<OwnerRRef>(iter->second);
}
}
RRefForkData RRefContext::prepareChildFork(const std::shared_ptr<RRef>& rref) {
auto rfd = rref->fork();
if (rref->isOwner()) {
// Note [Early Fork Registration]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// If the parent (caller) is the owner, directly register the fork, instead
// of waiting for another RREF_FORK_REQUEST or RREF_CHILD_ACCEPT message. An
// Alternative is adding the fork when the callee user ACKs. However, before
// that, the owner still have to adds the OwnerRRef into some map to keep it
// alive (e.g., in pendingChildren_). Hence, adding the fork here or in the
// ACK does not making any difference but only add complexity.
// TODO: When adding failure retries and timeout, this fork needs to be
// deleted if the owner does not receive the ACK within the timeout.
addForkOfOwner(rfd.rrefId_, rfd.forkId_);
// ensure that this RRef is in the owners_ list to keep it alive.
// this is needed for OwnerRRefs that were created locally.
{
std::lock_guard<std::mutex> lock(mutex_);
owners_[rref->rrefId()] = rref;
}
} else {
// Note [Useful Phantom Fork ID for User to Owner Call]
// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
// If the callee of dist.remote or dist.rpc is the owner of this RRef, the
// callee will not create a fork using this rfd.forkId_, because the owner
// will only keep one `OwnerRRef` instance and will not create any
// `UserRRef` instances. However, this rfd.forkId_ is still necessary, as
// the caller user needs to keep this `UserRRef` alive until it gets the
// ACK from the callee owner. Otherwise, the delete message could arrive
// at the owner before this dist.rpc or dist.remote call, which could
// potentially trigger the `OwnerRRef` to be deleted before running the
// user code.
addPendingChild(rfd.forkId_, rref);
}
return rfd;
}
void RRefContext::notifyOwnerAndParentOfFork(
const ForkId& forkId,
worker_id_t parent,
const std::shared_ptr<RRef>& rref) {
if (parent == rref->owner()) {
if (parent == agent_->getWorkerInfo().id_) {
// Owner sending RRef to self, remove the forkId as it was added during
// pickling
auto deletedRRef = delForkOfOwner(rref->rrefId(), forkId);
if (deletedRRef) {
TORCH_INTERNAL_ASSERT(
deletedRRef->rrefId() == rref->rrefId(),
"Deleting a fork of ",
rref->rrefId(),
" triggered deleting the OwnerRRef of ",
deletedRRef->rrefId());
// NB: not necessary to reset deletedRRef as rref is another shared_ptr
// instance pointing to the same OwnerRRef.
}
} else {
// If the parent is the owner, this fork has already been added into the
// forks_ map when the owner sends the message to the callee user. Hence,
// it is not necessary to send another RREF_CHILD_ACCEPT or
// RREF_FORK_REQUEST back to the owner. See Note [Early Fork
// Registration].
}
return;
}
if (rref->isOwner()) {
// See Note [Useful Phantom Fork ID for User to Owner Call]
// In this case, the owner is the caller, and it does not add the fork id
// into forks_. Because, there will be no real `UserRRef` associated with
// this fork ID.
auto fm = agent_->send(
agent_->getWorkerInfo(parent), RRefChildAccept(forkId).toMessage());
fm->addCallback([](const Message& /* unused */,
const c10::optional<utils::FutureError>& futErr) {
handleException(futErr);
});
} else {
auto fm = agent_->send(
agent_->getWorkerInfo(rref->owner()),
RRefForkRequest(rref->rrefId(), forkId).toMessage());
addPendingUser(forkId, rref);
fm->addCallback([this, forkId, parent](
const Message& /* unused */,
const c10::optional<utils::FutureError>& futErr) {
handleException(futErr);
this->finishForkRequest(forkId, parent);
});
}
}
void RRefContext::addPendingChild(
const ForkId& forkId,
const std::shared_ptr<RRef>& rref) {
// see Note [Early Fork Registration]
// If the parent is the owner, it should directly add the child UserRRef as a
// fork.
TORCH_INTERNAL_ASSERT(
!rref->isOwner(), "OwnerRRef should not have a pending child.");
std::lock_guard<std::mutex> lock(mutex_);
TORCH_INTERNAL_ASSERT(
pendingChildren_.find(forkId) == pendingChildren_.end(),
"Inconsistent states: attempt to add the same child fork twice.");
pendingChildren_[forkId] = rref;
}
void RRefContext::delPendingChild(const ForkId& forkId) {
std::lock_guard<std::mutex> lock(mutex_);
auto iter = pendingChildren_.find(forkId);
TORCH_INTERNAL_ASSERT(
iter != pendingChildren_.end(),
"Inconsistent states: attempt to delete a non-exist child fork.");
pendingChildren_.erase(iter);
}
void RRefContext::addPendingUser(
const ForkId& forkId,
const std::shared_ptr<RRef>& rref) {
std::lock_guard<std::mutex> lock(mutex_);
TORCH_INTERNAL_ASSERT(
pendingUsers_.find(forkId) == pendingUsers_.end(),
"Inconsistent states: attempt to add the same UserRRef twice.");
pendingUsers_[forkId] = rref;
}
void RRefContext::delPendingUser(const ForkId& forkId) {
std::lock_guard<std::mutex> lock(mutex_);
auto iter = pendingUsers_.find(forkId);
TORCH_INTERNAL_ASSERT(
iter != pendingUsers_.end(),
"Inconsistent states: attempt to delete a non-exist UserRRef.");
pendingUsers_.erase(iter);
}
void RRefContext::finishForkRequest(const ForkId& forkId, worker_id_t parent) {
delPendingUser(forkId);
auto fm = agent_->send(
agent_->getWorkerInfo(parent), RRefChildAccept(forkId).toMessage());
fm->addCallback([](const Message& /* unused */,
const c10::optional<utils::FutureError>& futErr) {
handleException(futErr);
});
}
void RRefContext::addSelfAsFork(std::shared_ptr<OwnerRRef>& rref) {
std::lock_guard<std::mutex> lock(mutex_);
const auto& rrefId = rref->rrefId();
owners_[rrefId] = rref;
auto& rrefForks = forks_[rrefId];
TORCH_INTERNAL_ASSERT(
rrefForks.find(rrefId) == rrefForks.end(),
"Attempt to add self as fork twice ",
rrefId);
rrefForks.insert(rrefId);
}
void RRefContext::addForkOfOwner(const RRefId& rrefId, const ForkId& forkId) {
std::lock_guard<std::mutex> lock(mutex_);
auto& rrefForks = forks_[rrefId];
TORCH_INTERNAL_ASSERT(
rrefForks.find(forkId) == rrefForks.end(),
"Got fork notification twice on the same RRef ",
forkId);
rrefForks.insert(forkId);
}
std::shared_ptr<RRef> RRefContext::delForkOfOwner(
const RRefId& rrefId,
const ForkId& forkId) {
std::shared_ptr<RRef> deletedRRef = nullptr;
{
std::lock_guard<std::mutex> lock(mutex_);
auto rrefIter = forks_.find(rrefId);
TORCH_INTERNAL_ASSERT(
rrefIter != forks_.end(),
"Inconsistent states, deleting a fork before the owner knows it.");
auto& rrefForks = rrefIter->second;
auto forkIter = rrefForks.find(forkId);
TORCH_INTERNAL_ASSERT(
forkIter != rrefForks.end(),
"Attempt to delete a non-exist fork ",
forkId);
rrefForks.erase(forkId);
if (rrefForks.empty()) {
auto ownerIter = owners_.find(rrefId);
if (ownerIter != owners_.end()) {
deletedRRef = ownerIter->second;
owners_.erase(ownerIter);
}
forks_.erase(rrefIter);
}
}
return deletedRRef;
}
} // namespace rpc
} // namespace distributed
} // namespace torch