diff --git a/ci/cudf_pandas_scripts/pandas-tests/run.sh b/ci/cudf_pandas_scripts/pandas-tests/run.sh index abde5e5d160..48ee4a05628 100755 --- a/ci/cudf_pandas_scripts/pandas-tests/run.sh +++ b/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 \ diff --git a/python/cudf/cudf/pandas/__init__.py b/python/cudf/cudf/pandas/__init__.py index ff445a63f74..bf88c950385 100644 --- a/python/cudf/cudf/pandas/__init__.py +++ b/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): diff --git a/python/cudf/cudf/pandas/__main__.py b/python/cudf/cudf/pandas/__main__.py index fb8569fa1d0..d4cb42d4c0b 100644 --- a/python/cudf/cudf/pandas/__main__.py +++ b/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() diff --git a/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh b/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh index a66f63c09b3..9c65b74d081 100755 --- a/python/cudf/cudf/pandas/scripts/run-pandas-tests.sh +++ b/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 \