Add ability to enable rmm pool on `cudf.pandas` import #15628

galipremsagar · 2024-05-01T15:12:15Z

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.

python/cudf/cudf/pandas/__init__.py

bdice · 2024-05-01T15:21:22Z

python/cudf/cudf/pandas/__init__.py

+        if cudf_pandas_mr is not None:
+            from rmm.mr import PoolMemoryResource
+
+            mr = PoolMemoryResource(


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

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

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.

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");

}

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

@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.

Do you think we need arena?

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.

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

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

@harrism @bdice This should be ready for review.

Co-authored-by: Bradley Dice <[email protected]>

harrism

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

harrism · 2024-05-01T22:54:13Z

python/cudf/cudf/pandas/__init__.py

+        if cudf_pandas_mr is not None:
+            from rmm.mr import PoolMemoryResource
+
+            mr = PoolMemoryResource(


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.

…option

galipremsagar · 2024-05-21T14:36:58Z

python/cudf/cudf/pandas/__init__.py

+
+    if (rmm_mode := os.getenv("CUDF_PANDAS_RMM_MODE", None)) is not None:
+        import rmm.mr
+        from rmm._lib.memory_resource import get_free_device_memory


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

madsbk

Looks good

python/cudf/cudf_pandas_tests/test_cudf_pandas.py

…#1567) This PR adds `get_free_device_memory` that returns free GPU memory necessary for rapidsai/cudf#15628 Authors: - GALI PREM SAGAR (https://github.com/galipremsagar) Approvers: - Mark Harris (https://github.com/harrism) - Bradley Dice (https://github.com/bdice) - Lawrence Mitchell (https://github.com/wence-) URL: #1567

Co-authored-by: Mads R. B. Kristensen <[email protected]>

bdice

A few suggestions/questions but none are blocking. Resolve as you see fit.

bdice · 2024-06-13T15:01:17Z

python/cudf/cudf/pandas/__init__.py

@@ -19,6 +22,46 @@ def install():
    loader = ModuleAccelerator.install("pandas", "cudf", "pandas")
    global LOADED
    LOADED = loader is not None
+    import os


Move this import to the top of the file.

bdice · 2024-06-13T15:24:29Z

python/cudf/cudf/pandas/__init__.py

+
+    if (rmm_mode := os.getenv("CUDF_PANDAS_RMM_MODE", None)) is not None:
+        import rmm.mr
+        from rmm.mr import available_device_memory


Should we import rmm at the top, too? I don’t think we have much benefit to lazy import here.

Also we should use the fully qualified name for rmm.mr.available_device_memory so we only have one import statement for rmm.mr.

bdice · 2024-06-13T15:27:31Z

python/cudf/cudf/pandas/__init__.py

+                initial_pool_size=free_memory,
+            )
+        else:
+            raise TypeError(f"Unsupported rmm mode: {rmm_mode}")


Suggested change

raise TypeError(f"Unsupported rmm mode: {rmm_mode}")

raise ValueError(f"Unsupported rmm mode: {rmm_mode}")

bdice · 2024-06-13T15:28:52Z

python/cudf/cudf/pandas/__init__.py

+            mr = rmm.mr.ManagedMemoryResource()
+            rmm.mr.set_current_device_resource(mr)
+        elif rmm_mode == "managed_pool":
+            rmm.reinitialize(


Why does this call reinitialize while other modes create and set the current device resource?

galipremsagar · 2024-06-13T16:54:51Z

/merge

@bdice

This PR addresses review comments made by @bdice here: #15628 (review) Authors: - GALI PREM SAGAR (https://github.com/galipremsagar) Approvers: - Bradley Dice (https://github.com/bdice) URL: #16021

Add ability to enable rmm pool on import

ed73c18

galipremsagar added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels May 1, 2024

galipremsagar requested review from harrism and beckernick May 1, 2024 15:12

galipremsagar self-assigned this May 1, 2024

galipremsagar requested a review from a team as a code owner May 1, 2024 15:12

galipremsagar requested review from vyasr and charlesbluca May 1, 2024 15:12

Merge branch 'branch-24.06' into enable_rmm_option

2dd9c3b

github-actions bot added the Python Affects Python cuDF API. label May 1, 2024

bdice requested changes May 1, 2024

View reviewed changes

python/cudf/cudf/pandas/__init__.py Outdated Show resolved Hide resolved

python/cudf/cudf/pandas/__init__.py Outdated Show resolved Hide resolved

python/cudf/cudf/pandas/__init__.py Outdated Show resolved Hide resolved

bdice reviewed May 1, 2024

View reviewed changes

Apply suggestions from code review

583c87f

Co-authored-by: Bradley Dice <[email protected]>

harrism requested changes May 1, 2024

View reviewed changes

galipremsagar added 2 commits May 21, 2024 02:55

Merge remote-tracking branch 'upstream/branch-24.06' into enable_rmm_…

8d1c367

…option

address reviews

1f2132b

galipremsagar requested review from harrism and bdice May 21, 2024 06:03

galipremsagar mentioned this pull request May 21, 2024

Add available_device_memory to fetch free amount of memory on a GPU rapidsai/rmm#1567

Merged

3 tasks

galipremsagar commented May 21, 2024

View reviewed changes

harrism approved these changes May 21, 2024

View reviewed changes

madsbk approved these changes May 22, 2024

View reviewed changes

python/cudf/cudf_pandas_tests/test_cudf_pandas.py Outdated Show resolved Hide resolved

Update python/cudf/cudf_pandas_tests/test_cudf_pandas.py

5ca4602

Co-authored-by: Mads R. B. Kristensen <[email protected]>

github-actions bot added the cuDF (Python) label Jun 12, 2024

galipremsagar changed the base branch from branch-24.06 to branch-24.08 June 12, 2024 14:22

update branch

7693533

github-actions bot added the cudf.pandas Issues specific to cudf.pandas label Jun 12, 2024

Fix percentage calculation

919fe0e

galipremsagar added 2 commits June 12, 2024 15:24

closest to 256

531a585

Merge branch 'branch-24.08' into enable_rmm_option

748e1f8

bdice approved these changes Jun 13, 2024

View reviewed changes

galipremsagar added the 5 - Ready to Merge Testing and reviews complete, ready to merge label Jun 13, 2024

rapids-bot bot merged commit 3cb3df3 into rapidsai:branch-24.08 Jun 13, 2024
75 checks passed

galipremsagar mentioned this pull request Jun 13, 2024

Refactor rmm usage in cudf.pandas #16021

Merged

3 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add ability to enable rmm pool on `cudf.pandas` import #15628

Add ability to enable rmm pool on `cudf.pandas` import #15628

galipremsagar commented May 1, 2024 •

edited

Loading

bdice May 1, 2024

galipremsagar May 1, 2024

harrism May 1, 2024

bdice May 2, 2024

galipremsagar May 2, 2024

madsbk May 6, 2024 •

edited

Loading

harrism May 6, 2024

madsbk May 7, 2024

harrism May 7, 2024

galipremsagar May 21, 2024

harrism left a comment

harrism May 1, 2024

galipremsagar May 21, 2024

madsbk left a comment

bdice left a comment

bdice Jun 13, 2024

bdice Jun 13, 2024

bdice Jun 13, 2024

bdice Jun 13, 2024

galipremsagar commented Jun 13, 2024

	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");
	}

	raise TypeError(f"Unsupported rmm mode: {rmm_mode}")
	raise ValueError(f"Unsupported rmm mode: {rmm_mode}")

Add ability to enable rmm pool on cudf.pandas import #15628

Add ability to enable rmm pool on cudf.pandas import #15628

Conversation

galipremsagar commented May 1, 2024 • edited Loading

Description

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

madsbk May 6, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

harrism left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

madsbk left a comment

Choose a reason for hiding this comment

bdice left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

galipremsagar commented Jun 13, 2024

Add ability to enable rmm pool on `cudf.pandas` import #15628

Add ability to enable rmm pool on `cudf.pandas` import #15628

galipremsagar commented May 1, 2024 •

edited

Loading

madsbk May 6, 2024 •

edited

Loading