API#

CuFile#

class kvikio.cufile.CuFile(file: Union[Path, str], flags: str = 'r')#

File handle for GPUDirect Storage (GDS)

close() → None#: Deregister the file and close the file

fileno() → int#: Get the file descriptor of the open file

open_flags() → int#: Get the flags of the file descriptor (see open(2))

pread(buf, size: Optional[int] = None, file_offset: int = 0, task_size: Optional[int] = None) → IOFuture#

Reads specified bytes from the file into device or host memory in parallel

pread reads the data from a specified file at a specified offset and size bytes into buf. The API works correctly for unaligned offsets and any data size, although the performance might not match the performance of aligned reads. See additional details in the notes below.

pread is non-blocking and returns a IOFuture that can be waited upon. It partitions the operation into tasks of size task_size for execution in the default thread pool.

Parameters#

buf: buffer-like or array-like: Device or host buffer to read into.
size: int, optional: Size in bytes to read.
file_offset: int, optional: Offset in the file to read from.
task_size: int, default=kvikio.defaults.task_size(): Size of each task in bytes.

Returns#

IOFuture: Future that on completion returns the size of bytes that were successfully read.

Notes#

KvikIO can only make use of GDS for reads that are aligned to a page boundary. For unaligned reads, KvikIO has to split the reads into aligned and unaligned parts. The GPU page size is 4kB, so all reads should be at an offset that is a multiple of 4096 bytes. If the desired file_offset is not a multiple of 4096, it is likely desirable to round down to the nearest multiple of 4096 and discard any undesired bytes from the resulting data. Similarly, it is optimal for size to be a multiple of 4096 bytes. When GDS isn’t used, this is less critical.

pwrite(buf, size: Optional[int] = None, file_offset: int = 0, task_size: Optional[int] = None) → IOFuture#

Writes specified bytes from device or host memory into the file in parallel

pwrite writes the data from buf to the file at a specified offset and size. The API works correctly for unaligned offset and data sizes, although the performance is not on-par with aligned writes. See additional details in the notes below.

pwrite is non-blocking and returns a IOFuture that can be waited upon. It partitions the operation into tasks of size task_size for execution in the default thread pool.

Parameters#

buf: buffer-like or array-like: Device or host buffer to write to.
size: int, optional: Size in bytes to write.
file_offset: int, optional: Offset in the file to write from.
task_size: int, default=kvikio.defaults.task_size(): Size of each task in bytes.

Returns#

IOFuture: Future that on completion returns the size of bytes that were successfully written.

Notes#

KvikIO can only make use of GDS for writes that are aligned to a page boundary. For unaligned writes, KvikIO has to split the writes into aligned and unaligned parts. The GPU page size is 4kB, so all writes should be at an offset that is a multiple of 4096 bytes. If the desired file_offset is not a multiple of 4096, it is likely desirable to round down to the nearest multiple of 4096 and discard any undesired bytes from the resulting data. Similarly, it is optimal for size to be a multiple of 4096 bytes. When GDS isn’t used, this is less critical.

raw_read(buf, size: Optional[int] = None, file_offset: int = 0, dev_offset: int = 0) → int#

Reads specified bytes from the file into the device memory

This is a low-level version of .read that doesn’t use threads and does not support host memory.

Parameters#

buf: buffer-like or array-like: Device buffer to read into.
size: int, optional: Size in bytes to read.
file_offset: int, optional: Offset in the file to read from.
dev_offset: int, optional: Offset in the buf to read from.

Returns#

int: The size of bytes that were successfully read.

Notes#

raw_write(buf, size: Optional[int] = None, file_offset: int = 0, dev_offset: int = 0) → int#

Writes specified bytes from the device memory into the file

This is a low-level version of .write that doesn’t use threads and does not support host memory.

Parameters#

buf: buffer-like or array-like: Device buffer to write to.
size: int, optional: Size in bytes to write.
file_offset: int, optional: Offset in the file to write from.
dev_offset: int, optional: Offset in the buf to write from.

Returns#

int: The size of bytes that were successfully written.

Notes#

read(buf, size: Optional[int] = None, file_offset: int = 0, task_size: Optional[int] = None) → int#

Reads specified bytes from the file into the device memory in parallel

This is a blocking version of .pread.

Parameters#

buf: buffer-like or array-like: Device buffer to read into.
size: int, optional: Size in bytes to read.
file_offset: int, optional: Offset in the file to read from.
task_size: int, default=kvikio.defaults.task_size(): Size of each task in bytes.

Returns#

int: The size of bytes that were successfully read.

Notes#

write(buf, size: Optional[int] = None, file_offset: int = 0, task_size: Optional[int] = None) → int#

Writes specified bytes from the device memory into the file in parallel

This is a blocking version of .pwrite.

Parameters#

buf: buffer-like or array-like: Device buffer to write to.
size: int, optional: Size in bytes to write.
file_offset: int, optional: Offset in the file to write from.
task_size: int, default=kvikio.defaults.task_size(): Size of each task in bytes.

Returns#

int: The size of bytes that were successfully written.

Notes#

class kvikio.cufile.IOFuture(handle)#

Future for CuFile IO

This class shouldn’t be used directly, instead non-blocking IO operations such as CuFile.pread and CuFile.pwrite returns an instance of this class. Use .get() to wait on the completion of the IO operation and retrieve the result.

done() → bool#

Return True if the future is done.

Returns#

bool: Whether the future is done or not

get() → int#

Retrieve the result of the IO operation that created this future

This call blocks until the IO operation finishes.

Returns#

int: The size of bytes that were read or written successfully.

Zarr#

class kvikio.zarr.GDSStore(*args, **kwargs)#

GPUDirect Storage (GDS) class using directories and files.

This class works like zarr.storage.DirectoryStore but implements getitems() in order to support direct reading into device memory. It uses KvikIO for reads and writes, which in turn will use GDS when applicable.

Notes#

GDSStore doesn’t implement _fromfile() thus non-array data such as meta data is always read into host memory. This is because only zarr.Array use getitems() to retrieve data.

getitems(keys: Sequence[str], *, contexts: Mapping[str, Mapping] = {}) → Mapping[str, Any]#

Retrieve data from multiple keys.

Parameters#

keysIterable[str]: The keys to retrieve
contexts: Mapping[str, Context]: A mapping of keys to their context. Each context is a mapping of store specific information. If the “meta_array” key exist, GDSStore use its values as the output array otherwise GDSStore.default_meta_array is used.

Returns#

Mapping: A collection mapping the input keys to their results.

Defaults#

kvikio.defaults.compat_mode() → bool#

Check if KvikIO is running in compatibility mode.

Notice, this is not the same as the compatibility mode in cuFile. That is, cuFile can run in compatibility mode while KvikIO is not.

When KvikIO is running in compatibility mode, it doesn’t load libcufile.so. Instead, reads and writes are done using POSIX.

Set the environment variable KVIKIO_COMPAT_MODE to enable/disable compatibility mode. By default, compatibility mode is enabled: - when libcufile cannot be found - when running in Windows Subsystem for Linux (WSL) - when /run/udev isn’t readable, which typically happens when running inside a docker image not launched with –volume /run/udev:/run/udev:ro

Return#

bool: Whether KvikIO is running in compatibility mode or not.

kvikio.defaults.compat_mode_reset(enable: bool) → None#

Reset the compatibility mode.

Use this function to enable/disable compatibility mode explicitly.

Parameters#

enablebool: Set to True to enable and False to disable compatibility mode

kvikio.defaults.get_num_threads() → int#

Get the number of threads of the thread pool.

Set the default value using num_threads_reset() or by setting the KVIKIO_NTHREADS environment variable. If not set, the default value is 1.

Return#

nthreads: int: The number of threads in the current thread pool.

kvikio.defaults.num_threads_reset(nthreads: int) → None#

Reset the number of threads in the default thread pool.

Waits for all currently running tasks to be completed, then destroys all threads in the pool and creates a new thread pool with the new number of threads. Any tasks that were waiting in the queue before the pool was reset will then be executed by the new threads. If the pool was paused before resetting it, the new pool will be paused as well.

Parameters#

nthreadsint: The number of threads to use.