From 55e849c9a1115ef34a87e436fca2dc6b4d524601 Mon Sep 17 00:00:00 2001 From: Tianyu Liu Date: Wed, 26 Feb 2025 11:07:20 -0500 Subject: [PATCH] Small fixes --- python/kvikio/kvikio/benchmarks/utils.py | 4 ++-- python/kvikio/kvikio/cufile_driver.py | 17 +++++++++-------- python/kvikio/kvikio/defaults.py | 16 +++++++++------- python/kvikio/kvikio/utils.py | 3 +-- python/kvikio/tests/test_cufile_driver.py | 18 +++++++++++------- python/kvikio/tests/test_defaults.py | 20 +++++++++++++------- 6 files changed, 45 insertions(+), 33 deletions(-) diff --git a/python/kvikio/kvikio/benchmarks/utils.py b/python/kvikio/kvikio/benchmarks/utils.py index fa25c361a4..02d97991de 100644 --- 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 index 59b9e1168f..8018415191 100644 --- 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 index 68e49891b4..3688be6a6e 100644 --- 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 index 6a03421d61..e79386023c 100644 --- 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 index 4602a8a2ba..fcc95c6cbc 100644 --- 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 index 098d2657e5..82c6327f5e 100644 --- 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(