API

CuFile

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

File handle for GPUDirect Storage (GDS)

close() None[source]

Deregister the file and close the file

fileno() int[source]

Get the file descriptor of the open file

open_flags() int[source]

Get the flags of the file descriptor (see open(2))

pread(buf, size: int | None = None, file_offset: int = 0, task_size: int | None = None) IOFuture[source]

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: int | None = None, file_offset: int = 0, task_size: int | None = None) IOFuture[source]

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: int | None = None, file_offset: int = 0, dev_offset: int = 0) int[source]

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

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.

raw_write(buf, size: int | None = None, file_offset: int = 0, dev_offset: int = 0) int[source]

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

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.

read(buf, size: int | None = None, file_offset: int = 0, task_size: int | None = None) int[source]

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

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.

write(buf, size: int | None = None, file_offset: int = 0, task_size: int | None = None) int[source]

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

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.

class kvikio.cufile.IOFuture(handle)[source]

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[source]

Return True if the future is done.

Returns:
bool

Whether the future is done or not

get() int[source]

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(path, normalize_keys=False, dimension_separator=None, *, compressor_config_overwrite: Mapping | None = None, decompressor_config_overwrite: Mapping | None = None)[source]

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.

Parameters:
pathstring

Location of directory to use as the root of the storage hierarchy.

normalize_keysbool, optional

If True, all store keys will be normalized to use lower case characters (e.g. ‘foo’ and ‘FOO’ will be treated as equivalent). This can be useful to avoid potential discrepancies between case-sensitive and case-insensitive file system. Default value is False.

dimension_separator{‘.’, ‘/’}, optional

Separator placed between the dimensions of a chunk.

compressor_config_overwrite

If not None, use this Mapping to specify what is written to the Zarr metadata file on disk (.zarray). Normally, Zarr writes the configuration[1] given by the compressor argument to the .zarray file. Use this argument to overwrite the normal configuration and use the specified Mapping instead.

decompressor_config_overwrite

If not None, use this Mapping to specify what compressor configuration[1] is used for decompressing no matter the configuration found in the Zarr metadata on disk (the .zarray file).

[1] https://github.com/zarr-developers/numcodecs/blob/cb155432/numcodecs/abc.py#L79

Notes

Atomic writes are used, which means that data are first written to a temporary file, then moved into place when the write is successfully completed. Files are only held open while they are being read or written and are closed immediately afterwards, so there is no need to manually close any files.

Safe to write in multiple threads or processes.

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

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[source]

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

kvikio.defaults.compat_mode_reset(enable: bool) None[source]

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[source]

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.

kvikio.defaults.num_threads_reset(nthreads: int) None[source]

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. The default value can be specified by setting the KVIKIO_NTHREADS environment variable. If not set, the default value is 1.