Add ability to enable rmm pool on cudf.pandas import

[galipremsagar]

Description

This PR enables allocating of rmm memory pool on cudf.pandas import using the following environment variables:

export CUDF_PANDAS_RMM_MODE="pool"

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.

[bdice]

Do we always want the user's specified MR to be wrapped in a pool?

[galipremsagar]

I don't know, open to suggestions here. Or we can change this as we know which works the best from experimenting.

[harrism]

Certainly not. In the example in the PR description you specify CudeAsyncMemoryResource. That already has a pool, you probably wouldn't always want to wrap that in an RMM PoolMemoryResource.

I thought in the meeting we discussed just providing a limited set of choices here. More sophisticated users should be able to import RMM first and configure the current device resource to override if they have very specific MR needs.

[bdice]

We do a similar thing in libcudf NVBench benchmarks. Here are the options we expose there, for reference:

cudf/cpp/benchmarks/fixture/nvbench_fixture.hpp

Lines 73 to 84 in 8cc4cc1

	inline std::shared_ptr<rmm::mr::device_memory_resource> create_memory_resource(
	std::string const& mode)
	{
	if (mode == "cuda") return make_cuda();
	if (mode == "pool") return make_pool();
	if (mode == "async") return make_async();
	if (mode == "arena") return make_arena();
	if (mode == "managed") return make_managed();
	if (mode == "managed_pool") return make_managed_pool();
	CUDF_FAIL("Unknown rmm_mode parameter: " + mode +
	"\nExpecting: cuda, pool, async, arena, managed, or managed_pool");
	}

[galipremsagar]

@harrism Do you feel comfortable having a similar option set for cudf.pandas?

[madsbk]

@harrism Do you feel comfortable having a similar option set for cudf.pandas?

I think this option set would be good, it covers the typical use cases.

[harrism]

Do you think we need arena?

[madsbk]

Do you think we need arena?

Maybe, I haven't tried it, but its fragmentation avoidance might be useful.
However, to limit the user complexity/documentation of cudf.pandas, I am also fine skipping arena. We can always add it later if it turns out to be useful.

[harrism]

I think you can also manually enable it by importing RMM and enabling it before importing cudf.pandas.

[galipremsagar]

I updated the PR to have simpler rmm mode selection and added pytests for the same.

@harrism @bdice This should be ready for review.

[harrism]

I don't think we should always wrap with a PoolMemoryResource

[harrism]

Certainly not. In the example in the PR description you specify CudeAsyncMemoryResource. That already has a pool, you probably wouldn't always want to wrap that in an RMM PoolMemoryResource.

I thought in the meeting we discussed just providing a limited set of choices here. More sophisticated users should be able to import RMM first and configure the current device resource to override if they have very specific MR needs.

[galipremsagar]

This method is being added in: rapidsai/rmm#1567

[madsbk]

Looks good

[madsbk]

Suggested change

	del env["CUDF_PANDAS_RMM_MODE"]

I don't you need this since env is a copy?