Benchmarks clean up

.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

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

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

python/kvikio/kvikio/benchmarks/__init__.py

@@ -0,0 +1,2 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.

python/kvikio/benchmarks/single-node-io.py → 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())
 

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