Prevent deadlock between RapidsBufferStore and RapidsBufferBase on close

[abellina]

Signed-off-by: Alessandro Bellina abellina@nvidia.com

Targeting this to 22.04 but I can target to 22.02 if desired.

Closes #4664.

Prevents a deadlock between RapidsBufferStore and RapidsBufferBase on shutdown, removing code that was added before that was causing #4664 (this patch no longer stops the context -> no longer causes shuffle id conflicts on register)

[abellina]

This PR: #4374, added a .close for the SparkSession in the mortgage suite, saying that there was a hang, but not explaining the root cause. I removed the .close and that uncovered a deadlock, handled here (I assume that this is the reason for the context close). I reverted my other change to remove all buffers because this keeps the SparkContext alive, and so the conflict of shuffle ids is no longer seen.

That said, I now see device memory leaks. I'll dig at that next. I am not sure if those are related to my change.

[abellina]

Ok the reason for the memory leaks is because of the shutdown hook logic in the JVM. The shutdown hooks, as far as I understand, run in parallel, so a registered hook for the MemoryCleaner runs at any random time, even before Spark is able to stop it's context, unregister shuffles, etc. One approach that would remove ambiguity would be to change the MemoryCleaner to not hook directly to the JVM, but instead provide the API so that a custom shutdown order could be imposed (potentially registering the MemoryCleaner against Spark's ShutdownHookManager).

I think this is a nice to have but it would make the code easier to follow. It took me a while to realize that the cuDF MemoryCleaner is racing against spark's ShutdownHookManager.

[abellina]

build

[jlowe]

Would be good to file an issue about the underlying problem that still isn't addressed, re: closing and restarting the Spark session

[abellina]

build

[abellina]

build

[abellina]

It is not clear to me why the previous one failed. I am looking to get some logs out of this to see if I can figure it out.

[abellina]

build

[abellina]

I am going to merge this. The previous time this failed, it seemed to die at GpuDeviceManagerSuite. Looking at the test, it seems to be a really simple test, that is just checking values constructed in the suite. So I am not sure why it failed, and I've had successful CI with this patch a few times.