driver open/close python bindings

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/cufile_driver.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import atexit
+
+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 must have a matching call to
+    `driver_close()`.
+
+    Normally, it is not required to open and close the cuFile driver since
+    it is done automatically.
+
+    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 must have a matching call to
+    `driver_close()`.
+
+    Raises
+    ------
+    RuntimeError
+        If cuFile isn't available.
+    """
+    return cufile_driver.driver_close()
+
+
+def initialize() -> None:
+    """Open the cuFile driver and close it again at module exit
+
+    Normally, it is not required to open and close the cuFile driver since
+    it is done automatically.
+
+    Raises
+    ------
+    RuntimeError
+        If cuFile isn't available.
+    """
+    driver_open()
+    atexit.register(driver_close)

python/kvikio/tests/test_cufile_driver.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
+# See file LICENSE for terms.
+
+import pytest
+
+import kvikio.cufile_driver
+import kvikio.defaults
+
+
+@pytest.mark.skipif(
+    kvikio.defaults.compat_mode(),
+    reason="cannot open the cuFile driver when already running in compatibility mode",
+)
+def test_open_and_close():
+    kvikio.cufile_driver.driver_open()
+    kvikio.cufile_driver.driver_close()