python bindings

cpp/include/kvikio/shim/cufile.hpp

@@ -53,7 +53,7 @@ class cuFileAPI {
   decltype(cuFileStreamDeregister)* StreamDeregister{nullptr};
 
  private:
-  // Don't call driver open and close directly, use
+  // Don't call driver open and close directly, use `.driver_open()` and `driver_close()`.
   decltype(cuFileDriverOpen)* _DriverOpen{nullptr};
   decltype(cuFileDriverClose)* _DriverClose{nullptr};
 
@@ -142,7 +142,7 @@ class cuFileAPI {
   }
 
   /**
-   * @brief Open cuFile driver
+   * @brief Open the cuFile driver
    *
    * cuFile accept multiple calls cufileDriverOpen(), only the first call opens the driver, but
    * every call should have a matching call to `.driver_close()`.
@@ -157,7 +157,7 @@ class cuFileAPI {
   }
 
   /**
-   * @brief Close cuFile driver
+   * @brief Close the cuFile driver
    */
   void driver_close()
   {

python/kvikio/kvikio/__init__.py

@@ -1,15 +1,10 @@
 # Copyright (c) 2021-2024, NVIDIA CORPORATION. All rights reserved.
 # See file LICENSE for terms.
 
-from kvikio._lib import cufile_driver  # type: ignore
 from kvikio._version import __git_commit__, __version__
 from kvikio.cufile import CuFile
 from kvikio.remote_file import RemoteFile, is_remote_file_available
 
-# TODO: Wrap nicely, maybe as a dataclass?
-DriverProperties = cufile_driver.DriverProperties
-
-
 __all__ = [
     "__git_commit__",
     "__version__",

python/kvikio/kvikio/_lib/cufile_driver.pyx

@@ -8,6 +8,19 @@
 from libcpp cimport bool
 
 
+cdef extern from "<kvikio/shim/cufile.hpp>" nogil:
+    cdef void cpp_driver_open "kvikio::cuFileAPI::instance().driver_open"() except +
+    cdef void cpp_driver_close "kvikio::cuFileAPI::instance().driver_close"() except +
+
+
+def driver_open():
+    cpp_driver_open()
+
+
+def driver_close():
+    cpp_driver_close()
+
+
 cdef extern from "<kvikio/cufile/driver.hpp>" nogil:
     cdef cppclass cpp_DriverProperties "kvikio::DriverProperties":
         cpp_DriverProperties() except +

python/kvikio/kvikio/benchmarks/utils.py

@@ -12,6 +12,7 @@
 from dask.utils import format_bytes
 
 import kvikio
+import kvikio.cufile_driver
 import kvikio.defaults
 
 
@@ -26,7 +27,7 @@ def drop_vm_cache() -> None:
 def pprint_sys_info() -> None:
     """Pretty print system information"""
 
-    props = kvikio.DriverProperties()
+    props = kvikio.cufile_driver.DriverProperties()
     try:
         import pynvml
 

python/kvikio/kvikio/cufile_driver.py

@@ -0,0 +1,37 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+from kvikio._lib import cufile_driver  # type: ignore
+
+# TODO: Wrap nicely, maybe as a dataclass?
+DriverProperties = cufile_driver.DriverProperties
+
+
+def driver_open() -> None:
+    """Open the cuFile driver
+
+    cuFile accept multiple calls to `driver_open`, only the first call
+    opens the driver, but every call should have a matching call to
+    `driver_close`.
+
+    Raises
+    ------
+    RuntimeError
+        If cuFile isn't available.
+    """
+    return cufile_driver.driver_open()
+
+
+def driver_close() -> None:
+    """Close the cuFile driver
+
+    cuFile accept multiple calls to `driver_open`, only the first call
+    opens the driver, but every call should have a matching call to
+    driver_close().
+
+    Raises
+    ------
+    RuntimeError
+        If cuFile isn't available.
+    """
+    return cufile_driver.driver_close()

python/kvikio/pyproject.toml

@@ -144,3 +144,6 @@ filterwarnings = [
     "ignore:Jitify is performing a one-time only warm-up to populate the persistent cache",
     "ignore::DeprecationWarning:botocore.*",
 ]
+markers = [
+    "cufile: tests to skip if cuFile isn't available e.g. run with `pytest -m 'not cufile'`"
+]

python/kvikio/tests/test_cufile_driver.py

@@ -0,0 +1,12 @@
+# Copyright (c) 2022-2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import pytest
+
+import kvikio.cufile_driver
+
+
+@pytest.mark.cufile
+def test_open_and_close():
+    kvikio.cufile_driver.driver_open()
+    kvikio.cufile_driver.driver_close()