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