Dask-cuda: Provide option for RMM logging

[jakirkham]

There's previously been interest in having a way to enable RMM logging in Dask workflows to help trackdown and identify memory allocation issues. This was worked on previously in PR ( https://github.com/rapidsai/dask-cuda/pull/322 ). Though the RMM API has changed some since that was opened. Would be good to revisit at some point and that functionality.

[charlesbluca]

I'm looking into this - currently have a working branch at:

• https://github.com/charlesbluca/dask-cuda/tree/rmm-logging

Could you go into some detail on what might need to be different to account for the API changes (especially the deprecation of CNMeM)? So far, my branch hasn't had to differ much from your work in the original PR.

[jakirkham]

Honestly Idk. There may not be any changes needed. Just noting that was based on an older version of RMM

[charlesbluca]

In that case, I'll just move ahead and use this issue as a sounding board if any problems come up with the newer version.

[charlesbluca]

The logging options seems to work, and when specified spawns an rmm.mr.LoggingResourceAdaptor for each worker. However, I need to rethink where the log files will be written, as worker.local_directory gets deleted when the client shuts down. Maybe in the current working directory under rmm_logs/<worker_name>.txt?

[jakirkham]

Yeah I think using the worker's local_directory makes sense. That's where the worker handles all other logging, spilling, and other things like this.

[jakirkham]

Oh wrote that too quickly. Yeah maybe we can take a parameter with this flag to specify a location?

[charlesbluca]

Done! Removed the Dask config fallback as I wasn't sure if we would want a default location for RMM logs to go - can always add it back if that's the case.

Is there a reliable way to produce RMM debug info for workers? Just need to check that the log files are being generated properly, and then maybe add a specific test case to check for the files.

[jakirkham]

Great, thank you! 😄

I think one of the benchmarks from the PTDS work should do the job. Main thing is to make sure the RMM pool is enabled

[charlesbluca]

Opened up a PR #542 - thanks for the help!

[charlesbluca]

PR is merged now - can we resolve this or should we leave it open for the possibility of scheduler/client logging?

[jakirkham]

Nope we can close it. Thanks for working on this Charles! 😄