-
Notifications
You must be signed in to change notification settings - Fork 96
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Update RMM tests based on deprecated CNMeM #359
Update RMM tests based on deprecated CNMeM #359
Conversation
As `get_default_resource_type` was dropped in RMM recently, use the newly introduce `rmm.mr.get_per_device_resource` instead to access the resource on device `0` (configured to a unique device for each worker). Since the memory resource itself is not realistically serializable, instead grab the type of each resource to send back. This is all done within a `lambda` to allow for a function that can be run on each worker.
As CNMeM has been dropped from Python and replaced with RMM's own pool resource, just check for that pool resource instead.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think there's a function for that. :)
Simplify the resource type checks a bit. Thanks Mark! :) Co-authored-by: Mark Harris <[email protected]>
Much cleaner. Thanks Mark! :) |
Codecov Report
@@ Coverage Diff @@
## branch-0.15 #359 +/- ##
===============================================
+ Coverage 59.65% 60.04% +0.39%
===============================================
Files 17 17
Lines 1321 1334 +13
===============================================
+ Hits 788 801 +13
Misses 533 533
Continue to review full report at Codecov.
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM, thanks @jakirkham !
The recent RMM PR ( rapidsai/rmm#466 ) deprecated CNMeM and made some changes for handling memory resources for multiple devices (needed by XGBoost). As a result, we are seeing a few test failures in Dask-CUDA. This makes the necessary changes to update these tests.