Repository: rapidsai/kvikio (language: python)

Issue:
Moving from `pynvml` to `nvidia-ml-py`
Currently this is making use of `pynvml` in a few places:

https://github.com/rapidsai/kvikio/blob/c152f633779d8c41791a3b76f9f678f926a9847d/python/benchmarks/single-node-io.py#L248-L250

However we would like to move to `nvidia-ml-py` in the future. Raising this issue to track this work


Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 788457c394..f794394103 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -40,8 +40,8 @@ repos:
         args: ["--config-file=python/kvikio/pyproject.toml",
                 "python/kvikio/kvikio",
                 "python/kvikio/tests",
-                "python/kvikio/examples",
-                "python/kvikio/benchmarks"]
+                "python/kvikio/examples"
+              ]
         pass_filenames: false
   - repo: https://github.com/pre-commit/mirrors-clang-format
     rev: v16.0.6
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 70f8279ef4..5dbbcb0ade 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -23,60 +23,6 @@ More information can be found at: [Contributor Code of Conduct](https://docs.rap
 
 ## Code contributions
 
-### Requirements
-
-To install users should have a working Linux machine with CUDA Toolkit
-installed (v11.4+) and a working compiler toolchain (C++17 and cmake).
-
-#### C++
-
-The C++ bindings are header-only and depends on CUDA Driver and Runtime API.
-In order to build and run the example code, CMake is required.
-
-#### Python
-
-The Python packages depends on the following packages:
-
-* Cython
-* Pip
-
-For testing:
-* pytest
-* cupy
-
-### Build KvikIO from source
-
-#### C++
-To build the C++ example, go to the `cpp` subdiretory and run:
-```
-mkdir build
-cd build
-cmake ..
-make
-```
-Then run the example:
-```
-./examples/basic_io
-```
-
-#### Python
-
-To build and install the extension, go to the `python` subdiretory and run:
-```
-python -m pip install .
-```
-One might have to define `CUDA_HOME` to the path to the CUDA installation.
-
-In order to test the installation, run the following:
-```
-pytest tests/
-```
-
-And to test performance, run the following:
-```
-python benchmarks/single-node-io.py
-```
-
 ### Code Formatting
 
 #### Using pre-commit hooks
diff --git a/docs/source/install.rst b/docs/source/install.rst
index 120329f81e..3e7985401c 100644
--- a/docs/source/install.rst
+++ b/docs/source/install.rst
@@ -61,11 +61,11 @@ In order to test the installation, run the following:
 
 .. code-block::
 
-  pytest tests/
+  pytest python/kvikio/tests/
 
 
 And to test performance, run the following:
 
 .. code-block::
 
-  python benchmarks/single-node-io.py
+  python python/kvikio/kvikio/benchmarks/single_node_io.py
diff --git a/python/kvikio/kvikio/benchmarks/__init__.py b/python/kvikio/kvikio/benchmarks/__init__.py
new file mode 100644
index 0000000000..8586c47db2
--- /dev/null
+++ b/python/kvikio/kvikio/benchmarks/__init__.py
@@ -0,0 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
diff --git a/python/kvikio/benchmarks/single-node-io.py b/python/kvikio/kvikio/benchmarks/single_node_io.py
similarity index 82%
rename from python/kvikio/benchmarks/single-node-io.py
rename to python/kvikio/kvikio/benchmarks/single_node_io.py
index 72b57300cc..4d47a80791 100644
--- a/python/kvikio/benchmarks/single-node-io.py
+++ b/python/kvikio/kvikio/benchmarks/single_node_io.py
@@ -1,10 +1,8 @@
-# Copyright (c) 2021-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 import argparse
 import contextlib
-import os
-import os.path
 import pathlib
 import shutil
 import statistics
@@ -17,6 +15,7 @@
 
 import kvikio
 import kvikio.defaults
+from kvikio.benchmarks.utils import parse_directory, pprint_sys_info
 
 
 def get_zarr_compressors() -> Dict[str, Any]:
@@ -260,53 +259,10 @@ def main(args):
     cupy.arange(10)  # Make sure CUDA is initialized
 
     kvikio.defaults.num_threads_reset(args.nthreads)
-    props = kvikio.DriverProperties()
-    try:
-        import pynvml.smi
-
-        nvsmi = pynvml.smi.nvidia_smi.getInstance()
-    except ImportError:
-        gpu_name = "Unknown (install pynvml)"
-        mem_total = gpu_name
-        bar1_total = gpu_name
-    else:
-        info = nvsmi.DeviceQuery()["gpu"][0]
-        gpu_name = f"{info['product_name']} (dev #0)"
-        mem_total = format_bytes(
-            parse_bytes(
-                str(info["fb_memory_usage"]["total"]) + info["fb_memory_usage"]["unit"]
-            )
-        )
-        bar1_total = format_bytes(
-            parse_bytes(
-                str(info["bar1_memory_usage"]["total"])
-                + info["bar1_memory_usage"]["unit"]
-            )
-        )
-    gds_version = "N/A (Compatibility Mode)"
-    if props.is_gds_available:
-        gds_version = f"v{props.major_version}.{props.minor_version}"
-    gds_config_json_path = os.path.realpath(
-        os.getenv("CUFILE_ENV_PATH_JSON", "/etc/cufile.json")
-    )
 
     print("Roundtrip benchmark")
     print("----------------------------------")
-    if kvikio.defaults.compat_mode():
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-        print("   WARNING - KvikIO compat mode   ")
-        print("      libcufile.so not used       ")
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-    elif not props.is_gds_available:
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-        print("   WARNING - cuFile compat mode   ")
-        print("         GDS not enabled          ")
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-    print(f"GPU               | {gpu_name}")
-    print(f"GPU Memory Total  | {mem_total}")
-    print(f"BAR1 Memory Total | {bar1_total}")
-    print(f"GDS driver        | {gds_version}")
-    print(f"GDS config.json   | {gds_config_json_path}")
+    pprint_sys_info()
     print("----------------------------------")
     print(f"nbytes            | {args.nbytes} bytes ({format_bytes(args.nbytes)})")
     print(f"4K aligned        | {args.nbytes % 4096 == 0}")
@@ -345,16 +301,6 @@ def pprint_api_res(name, samples):
 
 
 if __name__ == "__main__":
-
-    def parse_directory(x):
-        if x is None:
-            return x
-        else:
-            p = pathlib.Path(x)
-            if not p.is_dir():
-                raise argparse.ArgumentTypeError("Must be a directory")
-            return p
-
     parser = argparse.ArgumentParser(description="Roundtrip benchmark")
     parser.add_argument(
         "-n",
@@ -380,10 +326,10 @@ def parse_directory(x):
         help="Number of runs per API (default: %(default)s).",
     )
     parser.add_argument(
-        "--no-pre-register-buffer",
+        "--pre-register-buffer",
         action="store_true",
         default=False,
-        help="Disable pre-register of device buffer",
+        help="Enable pre-register of device buffer",
     )
     parser.add_argument(
         "-t",
@@ -413,7 +359,6 @@ def parse_directory(x):
     )
 
     args = parser.parse_args()
-    args.pre_register_buffer = args.no_pre_register_buffer is False
     if "all" in args.api:
         args.api = tuple(API.keys())
 
diff --git a/python/kvikio/kvikio/benchmarks/utils.py b/python/kvikio/kvikio/benchmarks/utils.py
new file mode 100644
index 0000000000..69375b8c21
--- /dev/null
+++ b/python/kvikio/kvikio/benchmarks/utils.py
@@ -0,0 +1,88 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+from __future__ import annotations
+
+import argparse
+import os
+import os.path
+import pathlib
+import subprocess
+
+from dask.utils import format_bytes
+
+import kvikio
+import kvikio.defaults
+
+
+def drop_vm_cache() -> None:
+    """Tells the Linux kernel to drop the page, inode, and dentry caches
+
+    See <https://linux-mm.org/Drop_Caches>
+    """
+    subprocess.check_output(["sudo /sbin/sysctl vm.drop_caches=3"], shell=True)
+
+
+def pprint_sys_info() -> None:
+    """Pretty print system information"""
+
+    props = kvikio.DriverProperties()
+    try:
+        import pynvml
+
+        pynvml.nvmlInit()
+        dev = pynvml.nvmlDeviceGetHandleByIndex(0)
+    except ImportError:
+        gpu_name = "Unknown (install nvidia-ml-py)"
+        mem_total = gpu_name
+        bar1_total = gpu_name
+    else:
+        gpu_name = f"{pynvml.nvmlDeviceGetName(dev)} (dev #0)"
+        mem_total = format_bytes(pynvml.nvmlDeviceGetMemoryInfo(dev).total)
+        bar1_total = format_bytes(pynvml.nvmlDeviceGetBAR1MemoryInfo(dev).bar1Total)
+    gds_version = "N/A (Compatibility Mode)"
+    if props.is_gds_available:
+        gds_version = f"v{props.major_version}.{props.minor_version}"
+    gds_config_json_path = os.path.realpath(
+        os.getenv("CUFILE_ENV_PATH_JSON", "/etc/cufile.json")
+    )
+
+    if kvikio.defaults.compat_mode():
+        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
+        print("   WARNING - KvikIO compat mode   ")
+        print("      libcufile.so not used       ")
+        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
+    elif not props.is_gds_available:
+        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
+        print("   WARNING - cuFile compat mode   ")
+        print("         GDS not enabled          ")
+        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
+    print(f"GPU               | {gpu_name}")
+    print(f"GPU Memory Total  | {mem_total}")
+    print(f"BAR1 Memory Total | {bar1_total}")
+    print(f"GDS driver        | {gds_version}")
+    print(f"GDS config.json   | {gds_config_json_path}")
+
+
+def parse_directory(x: str | None) -> pathlib.Path | None:
+    """Given an argparse argument, return a dir path.
+
+    None are passed through untouched.
+    Raise argparse.ArgumentTypeError if `x` isn't a directory (or None).
+
+    Parameters
+    ----------
+    x
+        argparse argument
+
+    Returns
+    -------
+    The directory path or None
+    """
+    if x is None:
+        return x
+    else:
+        p = pathlib.Path(x)
+        if not p.is_dir():
+            raise argparse.ArgumentTypeError("Must be a directory")
+        return p
diff --git a/python/kvikio/benchmarks/zarr-io.py b/python/kvikio/kvikio/benchmarks/zarr_io.py
similarity index 75%
rename from python/kvikio/benchmarks/zarr-io.py
rename to python/kvikio/kvikio/benchmarks/zarr_io.py
index 983c735364..fc226c2263 100644
--- a/python/kvikio/benchmarks/zarr-io.py
+++ b/python/kvikio/kvikio/benchmarks/zarr_io.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 import argparse
@@ -8,7 +8,6 @@
 import pathlib
 import shutil
 import statistics
-import subprocess
 import tempfile
 from time import perf_counter as clock
 from typing import ContextManager, Union
@@ -22,6 +21,7 @@
 import kvikio
 import kvikio.defaults
 import kvikio.zarr
+from kvikio.benchmarks.utils import drop_vm_cache, parse_directory, pprint_sys_info
 
 if not kvikio.zarr.supported:
     raise RuntimeError(f"requires Zarr >={kvikio.zarr.MINIMUM_ZARR_VERSION}")
@@ -32,11 +32,6 @@
 }
 
 
-def drop_vm_cache(args):
-    if args.drop_vm_cache:
-        subprocess.check_output(["sudo /sbin/sysctl vm.drop_caches=3"], shell=True)
-
-
 def create_src_data(args):
     return cupy.random.random(args.nelem, dtype=args.dtype)
 
@@ -51,7 +46,8 @@ def run_kvikio(args):
     src = create_src_data(args)
 
     # Write
-    drop_vm_cache(args)
+    if args.drop_vm_cache:
+        drop_vm_cache()
     t0 = clock()
     z = zarr.create(
         shape=(args.nelem,),
@@ -66,7 +62,8 @@ def run_kvikio(args):
     write_time = clock() - t0
 
     # Read
-    drop_vm_cache(args)
+    if args.drop_vm_cache:
+        drop_vm_cache()
     t0 = clock()
     res = z[:]
     read_time = clock() - t0
@@ -85,7 +82,8 @@ def run_posix(args):
     src = create_src_data(args)
 
     # Write
-    drop_vm_cache(args)
+    if args.drop_vm_cache:
+        drop_vm_cache()
     t0 = clock()
     z = zarr.create(
         shape=(args.nelem,),
@@ -100,7 +98,8 @@ def run_posix(args):
     write_time = clock() - t0
 
     # Read
-    drop_vm_cache(args)
+    if args.drop_vm_cache:
+        drop_vm_cache()
     t0 = clock()
     res = cupy.asarray(z[:])
     read_time = clock() - t0
@@ -120,57 +119,14 @@ def main(args):
     cupy.arange(10)  # Make sure CUDA is initialized
 
     kvikio.defaults.num_threads_reset(args.nthreads)
-    props = kvikio.DriverProperties()
-    try:
-        import pynvml.smi
-
-        nvsmi = pynvml.smi.nvidia_smi.getInstance()
-    except ImportError:
-        gpu_name = "Unknown (install pynvml)"
-        mem_total = gpu_name
-        bar1_total = gpu_name
-    else:
-        info = nvsmi.DeviceQuery()["gpu"][0]
-        gpu_name = f"{info['product_name']} (dev #0)"
-        mem_total = format_bytes(
-            parse_bytes(
-                str(info["fb_memory_usage"]["total"]) + info["fb_memory_usage"]["unit"]
-            )
-        )
-        bar1_total = format_bytes(
-            parse_bytes(
-                str(info["bar1_memory_usage"]["total"])
-                + info["bar1_memory_usage"]["unit"]
-            )
-        )
-    gds_version = "N/A (Compatibility Mode)"
-    if props.is_gds_available:
-        gds_version = f"v{props.major_version}.{props.minor_version}"
-    gds_config_json_path = os.path.realpath(
-        os.getenv("CUFILE_ENV_PATH_JSON", "/etc/cufile.json")
-    )
     drop_vm_cache_msg = str(args.drop_vm_cache)
     if not args.drop_vm_cache:
         drop_vm_cache_msg += " (use --drop-vm-cache for better accuracy!)"
     chunksize = args.chunksize * args.dtype.itemsize
 
-    print("Roundtrip benchmark")
+    print("Zarr-IO Benchmark")
     print("----------------------------------")
-    if kvikio.defaults.compat_mode():
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-        print("   WARNING - KvikIO compat mode   ")
-        print("      libcufile.so not used       ")
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-    elif not props.is_gds_available:
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-        print("   WARNING - cuFile compat mode   ")
-        print("         GDS not enabled          ")
-        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
-    print(f"GPU               | {gpu_name}")
-    print(f"GPU Memory Total  | {mem_total}")
-    print(f"BAR1 Memory Total | {bar1_total}")
-    print(f"GDS driver        | {gds_version}")
-    print(f"GDS config.json   | {gds_config_json_path}")
+    pprint_sys_info()
     print("----------------------------------")
     print(f"nbytes            | {args.nbytes} bytes ({format_bytes(args.nbytes)})")
     print(f"chunksize         | {chunksize} bytes ({format_bytes(chunksize)})")
@@ -211,16 +167,6 @@ def pprint_api_res(name, samples):
 
 
 if __name__ == "__main__":
-
-    def parse_directory(x):
-        if x is None:
-            return x
-        else:
-            p = pathlib.Path(x)
-            if not p.is_dir():
-                raise argparse.ArgumentTypeError("Must be a directory")
-            return p
-
     parser = argparse.ArgumentParser(description="Roundtrip benchmark")
     parser.add_argument(
         "-n",


Test patch:
diff --git a/python/kvikio/tests/test_benchmarks.py b/python/kvikio/tests/test_benchmarks.py
index ee0321d40a..3bdaf6613e 100644
--- a/python/kvikio/tests/test_benchmarks.py
+++ b/python/kvikio/tests/test_benchmarks.py
@@ -1,4 +1,4 @@
-# Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
+# Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
 import os
@@ -8,8 +8,9 @@
 
 import pytest
 
-benchmarks_path = Path(os.path.realpath(__file__)).parent / ".." / "benchmarks"
-
+benchmarks_path = (
+    Path(os.path.realpath(__file__)).parent.parent / "kvikio" / "benchmarks"
+)
 pytest.importorskip("cupy")
 pytest.importorskip("dask")
 
@@ -26,7 +27,7 @@
     ],
 )
 def test_single_node_io(run_cmd, tmp_path, api):
-    """Test benchmarks/single-node-io.py"""
+    """Test benchmarks/single_node_io.py"""
 
     if "zarr" in api:
         kz = pytest.importorskip("kvikio.zarr")
@@ -36,7 +37,37 @@ def test_single_node_io(run_cmd, tmp_path, api):
     retcode = run_cmd(
         cmd=[
             sys.executable or "python",
-            "single-node-io.py",
+            "single_node_io.py",
+            "-n",
+            "1MiB",
+            "-d",
+            str(tmp_path),
+            "--api",
+            api,
+        ],
+        cwd=benchmarks_path,
+    )
+    assert retcode == 0
+
+
+@pytest.mark.parametrize(
+    "api",
+    [
+        "kvikio",
+        "posix",
+    ],
+)
+def test_zarr_io(run_cmd, tmp_path, api):
+    """Test benchmarks/zarr_io.py"""
+
+    kz = pytest.importorskip("kvikio.zarr")
+    if not kz.supported:
+        pytest.skip(f"requires Zarr >={kz.MINIMUM_ZARR_VERSION}")
+
+    retcode = run_cmd(
+        cmd=[
+            sys.executable or "python",
+            "zarr_io.py",
             "-n",
             "1MiB",
             "-d",