Add ability to prefetch in cudf.pandas and change default to manage…

…d pool (#16296)

This PR adds ability to prefetch in `cudf.pandas` based off of: rapidsai/rmm#1608

Authors:
  - GALI PREM SAGAR (https://github.com/galipremsagar)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - Bradley Dice (https://github.com/bdice)
  - Muhammad Haseeb (https://github.com/mhaseeb123)
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Mark Harris (https://github.com/harrism)

URL: #16296

ci/cudf_pandas_scripts/pandas-tests/run.sh

@@ -19,7 +19,7 @@ RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${RESULTS_DIR}/test-results"}/
 mkdir -p "${RAPIDS_TESTS_DIR}"
 
 bash python/cudf/cudf/pandas/scripts/run-pandas-tests.sh \
-  -n 10 \
+  -n 5 \
   --tb=no \
   -m "not slow" \
   --max-worker-restart=3 \

python/cudf/cudf/pandas/__init__.py

@@ -25,41 +25,39 @@ def install():
     global LOADED
     LOADED = loader is not None
 
-    if (rmm_mode := os.getenv("CUDF_PANDAS_RMM_MODE", None)) is not None:
-        # Check if a non-default memory resource is set
-        current_mr = rmm.mr.get_current_device_resource()
-        if not isinstance(current_mr, rmm.mr.CudaMemoryResource):
-            warnings.warn(
-                f"cudf.pandas detected an already configured memory resource, ignoring 'CUDF_PANDAS_RMM_MODE'={str(rmm_mode)}",
-                UserWarning,
-            )
-        free_memory, _ = rmm.mr.available_device_memory()
-        free_memory = int(round(float(free_memory) * 0.80 / 256) * 256)
+    rmm_mode = os.getenv("CUDF_PANDAS_RMM_MODE", "managed_pool")
+    # Check if a non-default memory resource is set
+    current_mr = rmm.mr.get_current_device_resource()
+    if not isinstance(current_mr, rmm.mr.CudaMemoryResource):
+        warnings.warn(
+            f"cudf.pandas detected an already configured memory resource, ignoring 'CUDF_PANDAS_RMM_MODE'={str(rmm_mode)}",
+            UserWarning,
+        )
+        return rmm_mode
 
-        if rmm_mode == "cuda":
-            mr = rmm.mr.CudaMemoryResource()
-            rmm.mr.set_current_device_resource(mr)
-        elif rmm_mode == "pool":
-            rmm.mr.set_current_device_resource(
-                rmm.mr.PoolMemoryResource(
-                    rmm.mr.get_current_device_resource(),
-                    initial_pool_size=free_memory,
-                )
-            )
-        elif rmm_mode == "async":
-            mr = rmm.mr.CudaAsyncMemoryResource(initial_pool_size=free_memory)
-            rmm.mr.set_current_device_resource(mr)
-        elif rmm_mode == "managed":
-            mr = rmm.mr.ManagedMemoryResource()
-            rmm.mr.set_current_device_resource(mr)
-        elif rmm_mode == "managed_pool":
-            mr = rmm.mr.PoolMemoryResource(
+    free_memory, _ = rmm.mr.available_device_memory()
+    free_memory = int(round(float(free_memory) * 0.80 / 256) * 256)
+    new_mr = current_mr
+    if rmm_mode == "pool":
+        new_mr = rmm.mr.PoolMemoryResource(
+            current_mr,
+            initial_pool_size=free_memory,
+        )
+    elif rmm_mode == "async":
+        new_mr = rmm.mr.CudaAsyncMemoryResource(initial_pool_size=free_memory)
+    elif rmm_mode == "managed":
+        new_mr = rmm.mr.PrefetchResourceAdaptor(rmm.mr.ManagedMemoryResource())
+    elif rmm_mode == "managed_pool":
+        new_mr = rmm.mr.PrefetchResourceAdaptor(
+            rmm.mr.PoolMemoryResource(
                 rmm.mr.ManagedMemoryResource(),
                 initial_pool_size=free_memory,
             )
-            rmm.mr.set_current_device_resource(mr)
-        else:
-            raise ValueError(f"Unsupported rmm mode: {rmm_mode}")
+        )
+    elif rmm_mode != "cuda":
+        raise ValueError(f"Unsupported {rmm_mode=}")
+    rmm.mr.set_current_device_resource(new_mr)
+    return rmm_mode
 
 
 def pytest_load_initial_conftests(early_config, parser, args):

python/cudf/cudf/pandas/__main__.py

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION & AFFILIATES.
 # All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
@@ -72,7 +72,7 @@ def main():
 
     args = parser.parse_args()
 
-    install()
+    rmm_mode = install()
     with profile(args.profile, args.line_profile, args.args[0]) as fn:
         args.args[0] = fn
         if args.module:
@@ -86,6 +86,17 @@ def main():
             sys.argv[:] = args.args
             runpy.run_path(args.args[0], run_name="__main__")
 
+    if "managed" in rmm_mode:
+        for key in {
+            "column_view::get_data",
+            "mutable_column_view::get_data",
+            "gather",
+            "hash_join",
+        }:
+            from cudf._lib import pylibcudf
+
+            pylibcudf.experimental.enable_prefetching(key)
+
 
 if __name__ == "__main__":
     main()

python/cudf/cudf/pandas/scripts/run-pandas-tests.sh

@@ -137,7 +137,7 @@ and not test_eof_states \
 and not test_array_tz"
 
 # TODO: Remove "not db" once a postgres & mysql container is set up on the CI
-PANDAS_CI="1" timeout 30m python -m pytest -p cudf.pandas \
+PANDAS_CI="1" timeout 60m python -m pytest -p cudf.pandas \
     -v -m "not single_cpu and not db" \
     -k "$TEST_THAT_NEED_MOTO_SERVER and $TEST_THAT_CRASH_PYTEST_WORKERS and not test_groupby_raises_category_on_category and not test_constructor_no_pandas_array and not test_is_monotonic_na and not test_index_contains and not test_index_contains and not test_frame_op_subclass_nonclass_constructor and not test_round_trip_current" \
     --import-mode=importlib \