Skip to content

Commit 6580156

Browse files
tarang-jainfacebook-github-bot
authored andcommitted
Delete Raft Handle (#3435)
Summary: Small Raft related modification to StandardGpuResources: if the stream for a particular device is modified by a user, delete the Raft handle for that device. On any subsequent call to `getRaftHandle(device)`, a new raft handle with the updated stream will be created. Closes #3424 Pull Request resolved: #3435 Reviewed By: ramilbakhshyiev Differential Revision: D57640976 Pulled By: junjieqi fbshipit-source-id: 41e2898a39250b7e52e920b71e819fc21ca9fc85
1 parent eb28481 commit 6580156

File tree

1 file changed

+16
-0
lines changed

1 file changed

+16
-0
lines changed

faiss/gpu/StandardGpuResources.cpp

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,14 @@ void StandardGpuResourcesImpl::setDefaultStream(
257257
if (prevStream != stream) {
258258
streamWait({stream}, {prevStream});
259259
}
260+
#if defined USE_NVIDIA_RAFT
261+
// delete the raft handle for this device, which will be initialized
262+
// with the updated stream during any subsequent calls to getRaftHandle
263+
auto it2 = raftHandles_.find(device);
264+
if (it2 != raftHandles_.end()) {
265+
raftHandles_.erase(it2);
266+
}
267+
#endif
260268
}
261269

262270
userDefaultStreams_[device] = stream;
@@ -275,6 +283,14 @@ void StandardGpuResourcesImpl::revertDefaultStream(int device) {
275283

276284
streamWait({newStream}, {prevStream});
277285
}
286+
#if defined USE_NVIDIA_RAFT
287+
// delete the raft handle for this device, which will be initialized
288+
// with the updated stream during any subsequent calls to getRaftHandle
289+
auto it2 = raftHandles_.find(device);
290+
if (it2 != raftHandles_.end()) {
291+
raftHandles_.erase(it2);
292+
}
293+
#endif
278294
}
279295

280296
userDefaultStreams_.erase(device);

0 commit comments

Comments
 (0)