From 4a5884ce6ff94e4dfde2c0bdf8a6de068bb7de66 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Tue, 13 Aug 2024 08:46:55 +0200 Subject: [PATCH] benchmark: use cudf.set_option --- python/kvikio/kvikio/benchmarks/aws_s3_io.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/python/kvikio/kvikio/benchmarks/aws_s3_io.py b/python/kvikio/kvikio/benchmarks/aws_s3_io.py index 3c9d32e6c7..58e6222d07 100644 --- a/python/kvikio/kvikio/benchmarks/aws_s3_io.py +++ b/python/kvikio/kvikio/benchmarks/aws_s3_io.py @@ -96,9 +96,11 @@ def run() -> float: yield run() -def run_cudf(args, use_kvikio_s3): +def run_cudf(args, use_kvikio_s3: bool): import cudf + cudf.set_option("native_s3_io", use_kvikio_s3) + # Upload data to S3 server create_client_and_bucket() data = cupy.random.rand(args.nelem).astype(args.dtype) @@ -107,7 +109,7 @@ def run_cudf(args, use_kvikio_s3): def run() -> float: t0 = time.perf_counter() - cudf.read_parquet(f"s3://{args.bucket}/data1", use_kvikio_s3=use_kvikio_s3) + cudf.read_parquet(f"s3://{args.bucket}/data1") t1 = time.perf_counter() return t1 - t0