diff --git a/python/kvikio/kvikio/_lib/cufile_driver.pyx b/python/kvikio/kvikio/_lib/cufile_driver.pyx index 3028bfcbc6..29302a0104 100644 --- a/python/kvikio/kvikio/_lib/cufile_driver.pyx +++ b/python/kvikio/kvikio/_lib/cufile_driver.pyx @@ -8,6 +8,19 @@ from libcpp cimport bool +cdef extern from "" 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 "" nogil: cdef cppclass cpp_DriverProperties "kvikio::DriverProperties": cpp_DriverProperties() except + diff --git a/python/kvikio/kvikio/cufile_driver.py b/python/kvikio/kvikio/cufile_driver.py new file mode 100644 index 0000000000..8a0e625a55 --- /dev/null +++ b/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) diff --git a/python/kvikio/tests/test_cufile_driver.py b/python/kvikio/tests/test_cufile_driver.py new file mode 100644 index 0000000000..0980078205 --- /dev/null +++ b/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()