benchmark: use cudf.set_option

rapidsai · Aug 13, 2024 · 4a5884c · 4a5884c
1 parent 32130c2
commit 4a5884c
Showing 1 changed file with 4 additions and 2 deletions.
diff --git 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