diff --git a/benchmarks/cudf-merge.py b/benchmarks/cudf-merge.py index 982c3e93..da2093fb 100644 --- a/benchmarks/cudf-merge.py +++ b/benchmarks/cudf-merge.py @@ -160,7 +160,7 @@ async def worker(rank, eps, args): dev_id = args.devs[rank % len(args.devs)] cupy.cuda.runtime.setDevice(dev_id) rmm.reinitialize( - pool_allocator=True, devices=dev_id, initial_pool_size=args.rmm_init_pool_size + pool_allocator=True, devices=dev_id, initial_pool_size=args.rmm_init_pool_size, logging=True, log_file_name="cudf-merge-rmm.log", ) # Make cupy use RMM diff --git a/benchmarks/local-send-recv.py b/benchmarks/local-send-recv.py index 8d42342b..7c05f7cd 100644 --- a/benchmarks/local-send-recv.py +++ b/benchmarks/local-send-recv.py @@ -53,6 +53,8 @@ def server(queue, args): managed_memory=False, initial_pool_size=args.rmm_init_pool_size, devices=[args.server_dev], + logging=True, + log_file_name="server-local-send-recv-rmm.log", ) np.cuda.runtime.setDevice(args.server_dev) np.cuda.set_allocator(rmm.rmm_cupy_allocator) @@ -109,6 +111,8 @@ def client(queue, port, server_address, args): managed_memory=False, initial_pool_size=args.rmm_init_pool_size, devices=[args.client_dev], + logging=True, + log_file_name="client-local-send-recv-rmm.log", ) np.cuda.runtime.setDevice(args.client_dev) np.cuda.set_allocator(rmm.rmm_cupy_allocator)