From 0f07b0bb5e2cc89ca66e9d9639ff6ac961ec0471 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Tue, 30 Jul 2024 08:02:21 -0500 Subject: [PATCH] Enable prefetching before `runpy` (#16427) This PR enables prefetching before we execute the `runpy` module and script code. --- python/cudf/cudf/pandas/__main__.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/python/cudf/cudf/pandas/__main__.py b/python/cudf/cudf/pandas/__main__.py index d4cb42d4c0b..591744ce793 100644 --- a/python/cudf/cudf/pandas/__main__.py +++ b/python/cudf/cudf/pandas/__main__.py @@ -73,6 +73,16 @@ def main(): args = parser.parse_args() rmm_mode = install() + 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) with profile(args.profile, args.line_profile, args.args[0]) as fn: args.args[0] = fn if args.module: @@ -86,17 +96,6 @@ 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()