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 descripter of the open file

open_flags() int#

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

pread(buf, size: Optional[int] = None, file_offset: int = 0, task_size=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.

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.

IOFuture

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

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=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.

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.

IOFuture

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

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.

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.

int

The size of bytes that were successfully read.

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: 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.

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.

int

The size of bytes that were successfully written.

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

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

This is a blocking version of .pread.

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.

int

The size of bytes that were successfully read.

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

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

This is a blocking version of .pwrite.

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.

int

The size of bytes that were successfully written.

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)#

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.

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.

int

The size of bytes that were read or written successfully.

Zarr#

class kvikio.zarr.GDSStore(path, normalize_keys=False, dimension_separator=None)#

GPUDirect Storage (GDS) class using directories and files.

This class works like zarr.storage.DirectoryStore but use GPU buffers and will use GDS when applicable. The store supports both CPU and GPU buffers but when reading, GPU buffers are returned always.

TODO: Write metadata to disk in order to perserve the item types such that GPU items are read as GPU device buffers and CPU items are read as bytes.

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

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.

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 reset_num_threads() or by setting the KVIKIO_NTHREADS environment variable. If not set, the default value is 1.

nthreads: int

The number of threads in the current thread pool.

kvikio.defaults.reset_num_threads(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.

nthreadsint

The number of threads to use.