Open VibhuJawa opened 7 months ago
@quasiben , Wondering if we have opinion on this ? Happy to do a PR here to make life easier for folks like me.
CC: @alexbarghi-nv, @jnke2016 who have seen customer problems around a similar setup.
I have no objections to this, my only suggestion would be to make this a generic extensible option where we can then specify which libraries to set RMM as memory manager for, something like this:
--set-rmm-allocator=torch,another_future_library,...
Do you think that makes sense? @VibhuJawa if you want to get started on a PR for this I'm happy to help addressing any issues you may find along the way.
Currently we need to below to set rmm to use pytorch pool on a dask-cuda cluster. We should do this via a
cli