Small fixes

rapidsai · Feb 26, 2025 · 33b133e · 33b133e
1 parent d1bd127
commit 33b133e
Show file tree

Hide file tree

Showing 6 changed files with 45 additions and 33 deletions.
diff --git a/python/kvikio/kvikio/benchmarks/utils.py b/python/kvikio/kvikio/benchmarks/utils.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 from __future__ import annotations
@@ -28,7 +28,7 @@ def pprint_sys_info() -> None:
     """Pretty print system information"""
 
     version = kvikio.cufile_driver.libcufile_version()
-    props = kvikio.cufile_driver.DriverProperties()
+    props = kvikio.cufile_driver.properties
     try:
         import pynvml
 

diff --git a/python/kvikio/kvikio/cufile_driver.py b/python/kvikio/kvikio/cufile_driver.py
@@ -1,12 +1,11 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 import atexit
-from typing import Tuple, Any, overload
+from typing import Any, Tuple, overload
 
-from kvikio._lib import cufile_driver  # type: ignore
 import kvikio.utils
-
+from kvikio._lib import cufile_driver  # type: ignore
 
 properties = cufile_driver.DriverProperties()
 
@@ -46,10 +45,12 @@ def _set_property(self, property: str, value: Any):
     def _property_getter_and_setter(self) -> tuple[dict[str, Any], dict[str, Any]]:
         class_dict = vars(cufile_driver.DriverProperties)
 
-        property_getter_names = ["poll_mode",
-                                 "poll_thresh_size",
-                                 "max_device_cache_size",
-                                 "max_pinned_memory_size"]
+        property_getter_names = [
+            "poll_mode",
+            "poll_thresh_size",
+            "max_device_cache_size",
+            "max_pinned_memory_size",
+        ]
 
         property_getters = {}
         property_setters = {}

diff --git a/python/kvikio/kvikio/defaults.py b/python/kvikio/kvikio/defaults.py
@@ -43,13 +43,15 @@ def _set_property(self, property: str, value: Any):
     def _property_getter_and_setter(self) -> tuple[dict[str, Any], dict[str, Any]]:
         module_dict = vars(kvikio._lib.defaults)
 
-        property_getter_names = ["compat_mode",
-                                 "thread_pool_nthreads",
-                                 "task_size",
-                                 "gds_threshold",
-                                 "bounce_buffer_size",
-                                 "http_max_attempts",
-                                 "http_status_codes"]
+        property_getter_names = [
+            "compat_mode",
+            "thread_pool_nthreads",
+            "task_size",
+            "gds_threshold",
+            "bounce_buffer_size",
+            "http_max_attempts",
+            "http_status_codes",
+        ]
 
         property_getters = {}
         property_setters = {}

diff --git a/python/kvikio/kvikio/utils.py b/python/kvikio/kvikio/utils.py
@@ -79,8 +79,7 @@ def __enter__(self):
         else:
             handler = SimpleHTTPRequestHandler
 
-        handler_options = {**self.handler_options,
-                           **{"directory": self.root_path}}
+        handler_options = {**self.handler_options, **{"directory": self.root_path}}
 
         self.process = multiprocessing.Process(
             target=LocalHttpServer._server,

diff --git a/python/kvikio/tests/test_cufile_driver.py b/python/kvikio/tests/test_cufile_driver.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 import pytest
@@ -39,10 +39,14 @@ def test_property_setter():
 
     # Multiple context managers
     poll_mode_default = kvikio.cufile_driver.properties.poll_mode
-    max_device_cache_size_default = kvikio.cufile_driver.properties.max_device_cache_size
+    max_device_cache_size_default = (
+        kvikio.cufile_driver.properties.max_device_cache_size
+    )
     with kvikio.cufile_driver.set({"poll_mode": True, "max_device_cache_size": 2048}):
-        assert kvikio.cufile_driver.properties.poll_mode and\
-            (kvikio.cufile_driver.properties.max_device_cache_size == 2048)
-    assert (kvikio.cufile_driver.properties.poll_mode == poll_mode_default) and\
-        (kvikio.cufile_driver.properties.max_device_cache_size ==
-         max_device_cache_size_default)
+        assert kvikio.cufile_driver.properties.poll_mode and (
+            kvikio.cufile_driver.properties.max_device_cache_size == 2048
+        )
+    assert (kvikio.cufile_driver.properties.poll_mode == poll_mode_default) and (
+        kvikio.cufile_driver.properties.max_device_cache_size
+        == max_device_cache_size_default
+    )
diff --git a/python/kvikio/tests/test_defaults.py b/python/kvikio/tests/test_defaults.py
@@ -35,13 +35,19 @@ def test_property_setter():
     task_size_default = kvikio.defaults.task_size()
     num_threads_default = kvikio.defaults.num_threads()
     bounce_buffer_size_default = kvikio.defaults.bounce_buffer_size()
-    with kvikio.defaults.set({"task_size": 1024, "num_threads": 16, "bounce_buffer_size": 1024}):
-        assert (kvikio.defaults.task_size() == 1024) and\
-            (kvikio.defaults.num_threads() == 16) and\
-            (kvikio.defaults.bounce_buffer_size() == 1024)
-    assert (kvikio.defaults.task_size() == task_size_default) and\
-        (kvikio.defaults.num_threads() == num_threads_default) and\
-        (kvikio.defaults.bounce_buffer_size() == bounce_buffer_size_default)
+    with kvikio.defaults.set(
+        {"task_size": 1024, "num_threads": 16, "bounce_buffer_size": 1024}
+    ):
+        assert (
+            (kvikio.defaults.task_size() == 1024)
+            and (kvikio.defaults.num_threads() == 16)
+            and (kvikio.defaults.bounce_buffer_size() == 1024)
+        )
+    assert (
+        (kvikio.defaults.task_size() == task_size_default)
+        and (kvikio.defaults.num_threads() == num_threads_default)
+        and (kvikio.defaults.bounce_buffer_size() == bounce_buffer_size_default)
+    )
 
 
 @pytest.mark.skipif(