API#

CuFile#

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

File handle for GPUDirect Storage (GDS)

__init__(file: Path | str, flags: str = 'r')[source]#

Open and register file for GDS IO operations

CuFile opens the file twice and maintains two file descriptors. One file is opened with the specified flags and the other file is opened with the flags plus the O_DIRECT flag.

Parameters:
file: pathlib.Path or str

Path-like object giving the pathname (absolute or relative to the current working directory) of the file to be opened and registered.

flags: str, optional

“r” -> “open for reading (default)” “w” -> “open for writing, truncating the file first” “a” -> “open for writing, appending to the end of file if it exists” “+” -> “open for updating (reading and writing)”

close() None[source]#

Deregister the file and close the file

property closed: bool#
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.

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.

raw_read_async(
buf,
raw_stream: int,
size: int | None = None,
file_offset: int = 0,
dev_offset: int = 0,
) IOFutureStream[source]#

Reads specified bytes from the file into the device memory asynchronously

This is an async version of .raw_read that doesn’t use threads and does not support host memory.

Parameters:
buf: buffer-like or array-like

Device buffer to read into.

raw_stream: int

Raw CUDA stream to perform the read operation asynchronously.

size: int, optional

Size in bytes to read.

file_offset: int, optional

Offset in the file to read from.

Returns:
IOFutureStream

Future that when executed “.check_bytes_done()” returns the size of bytes that were successfully read. The instance must be kept alive until all data has been read from disk. One way to do this, is by calling IOFutureStream.check_bytes_done(), which will synchronize the associated stream and return the number of bytes read.

raw_write_async(
buf,
raw_stream: int,
size: int | None = None,
file_offset: int = 0,
dev_offset: int = 0,
) IOFutureStream[source]#

Writes specified bytes from the device memory into the file asynchronously

This is an async version of .raw_write that doesn’t use threads and does not support host memory.

Parameters:
buf: buffer-like or array-like

Device buffer to write to.

raw_stream: int

Raw CUDA stream to perform the write operation asynchronously.

size: int, optional

Size in bytes to write.

file_offset: int, optional

Offset in the file to write from.

Returns:
IOFutureStream

Future that when executed “.check_bytes_done()” returns the size of bytes that were successfully written. The instance must be kept alive until all data has been written to disk. One way to do this, is by calling IOFutureStream.check_bytes_done(), which will synchronize the associated stream and return the number of bytes written.

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.

is_direct_io_supported() bool[source]#

Whether Direct I/O is supported on this file handle.

This is determined by two factors: - Direct I/O support from the operating system and the file system - KvikIO global setting auto_direct_io_read and auto_direct_io_write. If both values are false, Direct I/O will not be supported on this file handle.

Returns:
bool

Whether Direct I/O is supported

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.

__init__(handle)[source]#
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.

done() bool[source]#

Return True if the future is done.

Returns:
bool

Whether the future is done or not

kvikio.cufile.get_page_cache_info(
file: PathLike | str | int | IOBase,
offset: int = 0,
length: int = 0,
) tuple[int, int][source]#

Obtain the page cache residency information for a given file

Example:

pages_cached, pages_total = kvikio.get_page_cache_info(my_file)
fraction_cached = pages_cached / pages_total
Parameters:
file: a path-like object, or string, or file descriptor, or file object

File to check.

offset: int, optional

Starting byte offset (default: 0 for beginning of file)

length: int, optional

Number of bytes to query (default: 0, meaning entire file from offset)

Returns:
tuple[int, int]

A pair containing the number of pages resident in the page cache and the total number of pages.

Notes

  • If offset is beyond the end of the file, returns (0, 0).

  • If offset + length extends beyond the file, the query is clamped to the file size.

  • The page cache residency query takes place in granularity of full pages. If the specified range does not align to page boundaries, partial pages at the start and end of the range are included.

kvikio.cufile.drop_file_page_cache(
file: PathLike | str | int | IOBase,
offset: int = 0,
length: int = 0,
sync_first: bool = True,
) None[source]#

Drop page cache for a specific file.

Advises the kernel to evict cached pages for the specified file descriptor using posix_fadvise with POSIX_FADV_DONTNEED.

Parameters:
file: a path-like object, or string, or file descriptor, or file object

File to operate on.

offset: int, optional

Starting byte offset (default: 0 for beginning of file)

length: int, optional

Number of bytes to drop (default: 0, meaning entire file from offset)

sync_first: bool, optional

Whether to flush dirty pages to disk before dropping. If True, fdatasync will be called prior to dropping. This ensures dirty pages become clean and thus droppable. Can be set to False if we are certain no dirty pages exist for this file.

Notes

  • This is the preferred method for benchmark cache invalidation as it:

    • Requires no elevated privileges

    • Affects only the specified file, not other processes

    • Has minimal overhead (no child process spawned)

  • The page cache dropping takes place in granularity of full pages. If the specified range does not align to page boundaries, partial pages at the start and end of the range are retained. Only pages fully contained within the range are dropped.

  • For dropping page cache system-wide (requires elevated privileges), see drop_system_page_cache().

kvikio.cufile.drop_system_page_cache(
reclaim_dentries_and_inodes: bool = True,
sync_first: bool = True,
) bool[source]#

Drop the system page cache.

Parameters:
reclaim_dentries_and_inodes: bool, optional

Whether to free reclaimable slab objects which include dentries and inodes.

  • If True, equivalent to executing /sbin/sysctl vm.drop_caches=3;

  • If False, equivalent to executing /sbin/sysctl vm.drop_caches=1.

sync_first: bool, optional

Whether to flush dirty pages to disk before dropping. If True, sync will be called prior to dropping. This ensures dirty pages become clean and thus droppable.

Returns:
bool

Whether the page cache has been successfully dropped

Notes

  • This drops page cache system-wide, affecting all processes. For dropping cache for a specific file without elevated privileges, see drop_file_page_cache().

  • This function creates a child process and executes the cache dropping shell command in the following order:

    • Execute the command without sudo prefix. This is for the superuser and also for specially configured systems where unprivileged users cannot execute /usr/bin/sudo but can execute /sbin/sysctl. If this step succeeds, the function returns True immediately.

    • Execute the command with sudo prefix. This is for the general case where selective unprivileged users have permission to run /sbin/sysctl with sudo prefix.

kvikio.cufile.clear_page_cache(*args, **kwargs)#

Deprecated since version 26.04: Use drop_system_page_cache() instead.

Drop the system page cache. Deprecated. Use drop_system_page_cache() instead.

kvikio.buffer.memory_register(buf) None[source]#

Register a device memory allocation with cuFile for GPUDirect Storage access.

This function automatically discovers the base address and size of the CUDA memory allocation containing buf. The entire underlying allocation is registered, regardless of which portion buf points to.

Registration pins the memory for GPU Direct DMA transfers, which can improve performance when the same buffer is reused across multiple cuFile I/O operations.

In compatibility mode (when GDS is unavailable), this function is a no-op.

Parameters:
buf: buffer-like or array-like

Device buffer to register .

kvikio.buffer.memory_deregister(buf) None[source]#

Deregister a device memory allocation from cuFile.

This function automatically discovers the base address of the CUDA memory allocation containing buf. The entire underlying allocation is deregistered, regardless of which portion buf points to.

In compatibility mode (when GDS is unavailable), this function is a no-op.

Parameters:
buf: buffer-like or array-like

Device buffer to deregister.

kvikio.buffer.bounce_buffer_free() int[source]#

Free the host allocations used as bounce buffers.

Returns:
Number of bytes freed.

CuFile driver#

class kvikio.cufile_driver.ConfigContextManager(config: dict[str, str])[source]#

Context manager allowing the cuFile driver configurations to be set upon entering a with block, and automatically reset upon leaving the block.

__init__(config: dict[str, str])[source]#
kvikio.cufile_driver.set(
config: dict[str, Any],
/,
) ConfigContextManager[source]#
kvikio.cufile_driver.set(
key: str,
value: Any,
/,
) ConfigContextManager

Set cuFile driver configurations.

Examples:

  • To set one or more properties

    # Set the property globally.
    kvikio.cufile_driver.set({"prop1": value1, "prop2": value2})
    
    # Set the property with a context manager.
    # The property automatically reverts to its old value
    # after leaving the `with` block.
    with kvikio.cufile_driver.set({"prop1": value1, "prop2": value2}):
        ...
    
  • To set a single property

    # Set the property globally.
    kvikio.cufile_driver.set("prop", value)
    
    # Set the property with a context manager.
    # The property automatically reverts to its old value
    # after leaving the `with` block.
    with kvikio.cufile_driver.set("prop", value):
        ...
    
Parameters:
config

The configurations. Can either be a single parameter (dict) consisting of one or more properties, or two parameters key (string) and value (Any) indicating a single property.

Valid configuration names are:

  • Read-only properties:

    • "is_gds_available"

    • "major_version"

    • "minor_version"

    • "allow_compat_mode"

    • "per_buffer_cache_size"

  • Settable properties:

    • "poll_mode"

    • "poll_thresh_size"

    • "max_device_cache_size"

    • "max_pinned_memory_size"

Returns:
ConfigContextManager

A context manager. If used in a with statement, the configuration will revert to its old value upon leaving the block.

kvikio.cufile_driver.get(config_name: str) Any[source]#

Get cuFile driver configurations.

Parameters:
config_name: str

The name of the configuration.

Valid configuration names are:

  • Read-only properties:

    • "is_gds_available"

    • "major_version"

    • "minor_version"

    • "allow_compat_mode"

    • "per_buffer_cache_size"

  • Settable properties:

    • "poll_mode"

    • "poll_thresh_size"

    • "max_device_cache_size"

    • "max_pinned_memory_size"

Returns:
Any

The value of the configuration.

kvikio.cufile_driver.libcufile_version() Tuple[int, int][source]#

Get the libcufile version.

Returns (0, 0) for cuFile versions prior to v1.8.

Returns:
The version as a tuple (MAJOR, MINOR).

Notes

This is not the version of the CUDA toolkit. cufile is part of the toolkit but follows its own version scheme.

kvikio.cufile_driver.driver_open() None[source]#

Open the cuFile driver

cuFile accepts 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.

kvikio.cufile_driver.driver_close() None[source]#

Close the cuFile driver

cuFile accepts 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.

kvikio.cufile_driver.initialize() None[source]#

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.

Notes

Registers an atexit handler that calls driver_close().

Mmap#

class kvikio.mmap.Mmap(
file_path: PathLike,
flags: str = 'r',
initial_map_size: int | None = None,
initial_map_offset: int = 0,
mode: int = 420,
map_flags: int | None = None,
)[source]#

Handle of a memory-mapped file

__init__(
file_path: PathLike,
flags: str = 'r',
initial_map_size: int | None = None,
initial_map_offset: int = 0,
mode: int = 420,
map_flags: int | None = None,
)[source]#

Construct a new memory-mapped file handle

Parameters:
file_pathos.PathLike

File path.

flagsstr, optional
  • r: Open for reading (default)

  • w: (Not implemented yet) Open for writing, truncating the file first

  • a: (Not implemented yet) Open for writing, appending to the end of file if it exists

  • +: (Not implemented yet) Open for updating (reading and writing)

initial_map_sizeint, optional

Size in bytes of the mapped region. If not specified, map the region starting from initial_map_offset to the end of file.

initial_map_offsetint, optional

File offset of the mapped region. Default is 0.

modeint, optional

Access mode (permissions) to use if creating a new file. Default is 0644 (octal), 420 (decimal).

map_flagsint, optional

Flags to be passed to the system call mmap. See mmap(2) for details.

initial_map_size() int[source]#

Size in bytes of the mapped region when the mapping handle was constructed

Returns:
int

Initial size of the mapped region.

initial_map_offset() int[source]#

File offset of the mapped region when the mapping handle was constructed

Returns:
int

Initial file offset of the mapped region.

file_size() int[source]#

Get the file size if the file is open

Returns 0 if the file is closed.

Returns:
int

The file size in bytes.

close() None[source]#

Close the mapping handle if it is open; do nothing otherwise

Unmaps the memory region and closes the underlying file descriptor.

closed() bool[source]#

Whether the mapping handle is closed

Returns:
bool

Boolean answer.

read(
buf: Any,
size: int | None = None,
offset: int = 0,
) int[source]#

Sequential read size bytes from the file to the destination buffer buf

Parameters:
bufbuffer-like or array-like

Address of the host or device memory (destination buffer).

sizeint, optional

Size in bytes to read. If not specified, read starts from offset to the end of file.

offsetint, optional

File offset. Default is 0.

Returns:
int

Number of bytes that have been read.

Raises:
IndexError

If the read region specified by offset and size is outside the initial region specified when the mapping handle was constructed.

RuntimeError

If the mapping handle is closed.

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

Parallel read size bytes from the file to the destination buffer buf

Parameters:
bufbuffer-like or array-like

Address of the host or device memory (destination buffer).

sizeint, optional

Size in bytes to read. If not specified, read starts from offset to the end of file.

offsetint, optional

File offset. Default is 0.

task_sizeint, optional

Size of each task in bytes for parallel execution. If None, uses the default task size from kvikio.defaults.task_size().

Returns:
IOFuture

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

Raises:
IndexError

If the read region specified by offset and size is outside the initial region specified when the mapping handle was constructed.

RuntimeError

If the mapping handle is closed.

Notes

The returned IOFuture object’s get() should not be called after the lifetime of the MmapHandle object ends. Otherwise, the behavior is undefined.

Zarr#

class kvikio.zarr.GDSStore(root: Path | str, *, read_only: bool = False)[source]#
async get(
key: str,
prototype: BufferPrototype | None = None,
byte_range: RangeByteRequest | OffsetByteRequest | SuffixByteRequest | None = None,
) Buffer | None[source]#

Retrieve the value associated with a given key.

Parameters:
keystr
prototypeBufferPrototype

The prototype of the output buffer. Stores may support a default buffer prototype.

byte_rangeByteRequest, optional

ByteRequest may be one of the following. If not provided, all data associated with the key is retrieved. - RangeByteRequest(int, int): Request a specific range of bytes in the form (start, end). The end is exclusive. If the given range is zero-length or starts after the end of the object, an error will be returned. Additionally, if the range ends after the end of the object, the entire remainder of the object will be returned. Otherwise, the exact requested range will be returned. - OffsetByteRequest(int): Request all bytes starting from a given byte offset. This is equivalent to bytes={int}- as an HTTP header. - SuffixByteRequest(int): Request the last int bytes. Note that here, int is the size of the request, not the byte offset. This is equivalent to bytes=-{int} as an HTTP header.

Returns:
Buffer
async set(key: str, value: Buffer) None[source]#

Store a (key, value) pair.

Parameters:
keystr
valueBuffer

RemoteFile#

kvikio.remote_file.infer_remote_endpoint_type(
url: str,
) RemoteEndpointType[source]#

Infer endpoint type from URL using AUTO endpoint resolution rules.

class kvikio.remote_file.RemoteEndpointType(*values)[source]#

Types of remote file endpoints supported by KvikIO.

This enum defines the different protocols and services that can be used to access remote files. It is used to specify or detect the type of remote endpoint when opening files.

Attributes:
AUTOint

Automatically detect the endpoint type from the URL. KvikIO will attempt to infer the appropriate protocol based on the URL format.

S3int

AWS S3 endpoint using credentials-based authentication. Requires AWS environment variables (such as AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, AWS_DEFAULT_REGION) to be set.

S3_PUBLICINT

AWS S3 endpoint for publicly accessible objects. No credentials required as the objects have public read permissions enabled. Used for open datasets and public buckets.

S3_PRESIGNED_URLint

AWS S3 endpoint using a presigned URL. No credentials required as authentication is embedded in the URL with time-limited access.

WEBHDFSint

Apache Hadoop WebHDFS (Web-based Hadoop Distributed File System) endpoint for accessing files stored in HDFS over HTTP/HTTPS.

HTTPint

Generic HTTP or HTTPS endpoint for accessing files from web servers. This is used for standard web resources that do not fit the other specific categories.

See also

RemoteFile.open

Factory method that uses this enum to specify endpoint types.

AUTO = 0#
S3 = 1#
S3_PUBLIC = 2#
S3_PRESIGNED_URL = 3#
WEBHDFS = 4#
HTTP = 5#
class kvikio.remote_file.RemoteFile(handle)[source]#

File handle of a remote file.

__init__(handle)[source]#

Create a remote file from a Cython handle.

This constructor should not be called directly instead use a factory method like RemoteFile.open_http()

Parameters:
handlekvikio._lib.remote_handle.RemoteFile

The Cython handle

classmethod open_http(
url: str,
nbytes: int | None = None,
) RemoteFile[source]#

Open a HTTP/HTTPS file.

Parameters:
url

URL to the remote file.

nbytes

The size of the file. If None, KvikIO will ask the server for the file size.

classmethod open_s3(
bucket_name: str,
object_name: str,
nbytes: int | None = None,
aws_region_name: str | None = None,
aws_access_key_id: str | None = None,
aws_secret_access_key: str | None = None,
aws_endpoint_url: str | None = None,
aws_session_token: str | None = None,
) RemoteFile[source]#

Open a AWS S3 file from a bucket name and object name.

AWS credentials can be provided as keyword arguments or through environment variables:

  • AWS_DEFAULT_REGION (or region_name parameter)

  • AWS_ACCESS_KEY_ID (or access_key_id parameter)

  • AWS_SECRET_ACCESS_KEY (or secret_access_key parameter)

  • AWS_SESSION_TOKEN (or aws_session_token parameter, when using temporary credentials)

Additionally, to overwrite the AWS endpoint, set AWS_ENDPOINT_URL (or endpoint_url parameter). See <https://docs.aws.amazon.com/cli/v1/userguide/cli-configure-envvars.html>

Parameters:
bucket_name

The bucket name of the file.

object_name

The object name of the file.

nbytes

The size of the file. If None, KvikIO will ask the server for the file size.

aws_region

The AWS region, such as “us-east-1”, to use. If None, the value of the AWS_DEFAULT_REGION environment variable is used.

aws_access_key

The AWS access key to use. If None, the value of the AWS_ACCESS_KEY_ID environment variable is used.

aws_secret_access_key

The AWS secret access key to use. If None, the value of the AWS_SECRET_ACCESS_KEY environment variable is used.

aws_endpoint_url

Overwrite the endpoint url (including the protocol part) by using the scheme: “<aws_endpoint_url>/<bucket_name>/<object_name>”. If None, the value of the AWS_ENDPOINT_URL environment variable is used. If this is also not set, the regular AWS url scheme is used: “https://<bucket_name>.s3.<region>.amazonaws.com/<object_name>”.

aws_session_token

The AWS session token to use. If None, the value of the AWS_SESSION_TOKEN environment variable is used.

classmethod open_s3_url(
url: str,
nbytes: int | None = None,
aws_region_name: str | None = None,
aws_access_key_id: str | None = None,
aws_secret_access_key: str | None = None,
aws_endpoint_url: str | None = None,
aws_session_token: str | None = None,
) RemoteFile[source]#

Open a AWS S3 file from an URL.

The url can take two forms:

AWS credentials can be provided as keyword arguments or through environment variables:

  • AWS_DEFAULT_REGION (or region_name parameter)

  • AWS_ACCESS_KEY_ID (or access_key_id parameter)

  • AWS_SECRET_ACCESS_KEY (or secret_access_key parameter)

  • AWS_SESSION_TOKEN (or aws_session_token parameter, when using temporary credentials)

Additionally, if url is a S3 url, it is possible to overwrite the AWS endpoint by setting AWS_ENDPOINT_URL (or endpoint_url parameter). See <https://docs.aws.amazon.com/cli/v1/userguide/cli-configure-envvars.html>

Parameters:
url

Either a http url or a S3 url.

nbytes

The size of the file. If None, KvikIO will ask the server for the file size.

aws_region

The AWS region, such as “us-east-1”, to use. If None, the value of the AWS_DEFAULT_REGION environment variable is used.

aws_access_key

The AWS access key to use. If None, the value of the AWS_ACCESS_KEY_ID environment variable is used.

aws_secret_access_key

The AWS secret access key to use. If None, the value of the AWS_SECRET_ACCESS_KEY environment variable is used.

aws_endpoint_url

Overwrite the endpoint url (including the protocol part) by using the scheme: “<aws_endpoint_url>/<bucket_name>/<object_name>”. If None, the value of the AWS_ENDPOINT_URL environment variable is used. If this is also not set, the regular AWS url scheme is used: “https://<bucket_name>.s3.<region>.amazonaws.com/<object_name>”.

aws_session_token

The AWS session token to use. If None, the value of the AWS_SESSION_TOKEN environment variable is used.

classmethod open_s3_public(
url: str,
nbytes: int | None = None,
) RemoteFile[source]#

Open a publicly accessible AWS S3 file.

Parameters:
url

URL to the remote file.

nbytes

The size of the file. If None, KvikIO will ask the server for the file size.

classmethod open_s3_presigned_url(
presigned_url: str,
nbytes: int | None = None,
) RemoteFile[source]#

Open a AWS S3 file from a presigned URL.

Parameters:
presigned_url

Presigned URL to the remote file.

nbytes

The size of the file. If None, KvikIO will ask the server for the file size.

classmethod open_webhdfs(
url: str,
nbytes: int | None = None,
) RemoteFile[source]#

Open a file on Apache Hadoop Distributed File System (HDFS) using WebHDFS.

If KvikIO is run within a Docker, the argument --network host needs to be passed to the docker run command.

Parameters:
url

URL to the remote file.

nbytes

The size of the file. If None, KvikIO will ask the server for the file size.

classmethod open(
url: str,
remote_endpoint_type: RemoteEndpointType = RemoteEndpointType.AUTO,
allow_list: list | None = None,
nbytes: int | None = None,
) RemoteFile[source]#

Create a remote file handle from a URL.

This function creates a RemoteFile for reading data from various remote endpoints including HTTP/HTTPS servers, AWS S3 buckets, S3 for public access, S3 presigned URLs, and WebHDFS. The endpoint type can be automatically detected from the URL or explicitly specified.

Parameters:
urlstr

The URL of the remote file. Supported formats include:

  • S3 with credentials

  • S3 for public access

  • S3 presigned URL

  • WebHDFS

  • HTTP/HTTPS

remote_endpoint_typeRemoteEndpointType, optional

The type of remote endpoint. Default is RemoteEndpointType.AUTO which automatically detects the endpoint type from the URL. Can be explicitly set to RemoteEndpointType.S3, RemoteEndpointType.S3_PUBLIC, RemoteEndpointType.S3_PRESIGNED_URL, RemoteEndpointType.WEBHDFS, or RemoteEndpointType.HTTP to force a specific endpoint type.

allow_listlist of RemoteEndpointType, optional

List of allowed endpoint types. If provided:

  • If remote_endpoint_type is RemoteEndpointType.AUTO, types are tried in the exact order specified until a match is found.

  • In explicit mode, the specified type must be in this list, otherwise an exception is thrown.

If not provided, defaults to all supported types in this order: RemoteEndpointType.S3, RemoteEndpointType.S3_PUBLIC, RemoteEndpointType.S3_PRESIGNED_URL, RemoteEndpointType.WEBHDFS, and RemoteEndpointType.HTTP.

nbytesint, optional

File size in bytes. If not provided, the function sends an additional request to the server to query the file size.

Returns:
RemoteFile

A RemoteFile object that can be used to read data from the remote file.

Raises:
RuntimeError
  • If the URL is malformed or missing required components.

  • RemoteEndpointType.AUTO mode is used and the URL does not match any supported endpoint type.

  • The specified endpoint type is not in the allow_list.

  • The URL is invalid for the specified endpoint type.

  • Unable to connect to the remote server or determine file size (when nbytes not provided).

Examples

  • Auto-detect endpoint type from URL:

    handle = RemoteFile.open(
        "https://bucket.s3.amazonaws.com/object?X-Amz-Algorithm=AWS4-HMAC-SHA256"
        "&X-Amz-Credential=...&X-Amz-Signature=..."
    )
    
  • Open S3 file with explicit endpoint type:

    handle = RemoteFile.open(
        "https://my-bucket.s3.us-east-1.amazonaws.com/data.bin",
        remote_endpoint_type=RemoteEndpointType.S3
    )
    
  • Restrict endpoint type candidates:

    handle = RemoteFile.open(
        user_provided_url,
        remote_endpoint_type=RemoteEndpointType.AUTO,
        allow_list=[
            RemoteEndpointType.HTTP,
            RemoteEndpointType.S3_PRESIGNED_URL
        ]
    )
    
  • Provide known file size to skip HEAD request:

    handle = RemoteFile.open(
        "https://example.com/large-file.bin",
        remote_endpoint_type=RemoteEndpointType.HTTP,
        nbytes=1024 * 1024 * 100  # 100 MB
    )
    
close() None[source]#

Close the file

remote_endpoint_type() RemoteEndpointType[source]#

Get the type of the remote file.

Returns:
The type of the remote file.
nbytes() int[source]#

Get the file size.

Note, this is very fast, no communication needed.

Returns:
The number of bytes.
read(
buf,
size: int | None = None,
file_offset: int = 0,
) int[source]#

Read from remote source into buffer (host or device memory) in parallel.

Parameters:
bufbuffer-like or array-like

Device or host buffer to read into.

size

Size in bytes to read.

file_offset

Offset in the file to read from.

Returns:
The size of bytes that were successfully read.
pread(
buf,
size: int | None = None,
file_offset: int = 0,
) IOFuture[source]#

Read from remote source into buffer (host or device memory) in parallel.

Parameters:
bufbuffer-like or array-like

Device or host buffer to read into.

size

Size in bytes to read.

file_offset

Offset in the file to read from.

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

Statistics#

class kvikio.statistics.SummaryMonitor[source]#

Turns on I/O statistics for the process and accumulates them while it exists

The intended use is to create one early, keep it, and read it whenever a report is wanted:

monitor = kvikio.SummaryMonitor()
...
print(monitor.get())

Or scope it to a phase, and ask for the interval:

with kvikio.SummaryMonitor() as monitor:
    before = monitor.get()
    run_a_phase()
    print(monitor.since(before).busy_bytes_per_sec)

KvikIO does no counting at all while no monitor exists. Counting happens entirely in C++, where the monitor is told when each operation starts and finishes, so Python pays only when a reading is taken, not per operation.

Notes

A monitor measures the whole process, not a scope. It counts every thread’s I/O while it exists and cannot attribute I/O to a particular call, so wrapping a block in one measures that block only if nothing else is doing I/O at the same time.

Monitors are independent: any number can exist at once, nested or overlapping, and resetting one has no effect on the others.

__init__()[source]#

Create a monitor and begin counting

get() Summary[source]#

Read the totals accumulated since construction, or since the last reset

Safe to call repeatedly, and non-destructive.

Returns:
Summary

The totals.

reset() None[source]#

Zero the totals and restart the span, as if the monitor had just been created

since(
previous: Summary,
) Summary[source]#

Totals for the interval since an earlier reading

Parameters:
previous

An earlier reading from this monitor.

Returns:
Summary

The interval’s totals, spanning [previous.end_unix_ns, now).

Raises:
ValueError

If previous is not an earlier reading of this monitor’s current span. See Summary.since().

stop() None[source]#

Stop counting. Idempotent, and one-way, as there is no resuming

The end of the measured span is fixed here, so later readings describe the interval that was measured rather than growing with the process.

class kvikio.statistics.Summary(
start_unix_ns: int,
end_unix_ns: int,
num_ops: int,
num_reads: int,
num_writes: int,
bytes_requested: int,
bytes_transferred: int,
bytes_read: int,
bytes_written: int,
num_errors: int,
busy_ns: int,
total_duration_ns: int,
by_backend: dict[str, BackendTotals],
counters: dict[str, int],
wall_ns: int,
busy_bytes_per_sec: float,
busy_fraction: float,
mean_duration_ns: int,
)[source]#

Totals of the I/O KvikIO has performed

A snapshot taken from a SummaryMonitor, whose values do not change once read. This class shouldn’t be constructed directly, use SummaryMonitor.get() or SummaryMonitor.since().

Everything here describes logical operations: one read() is one operation however many reads KvikIO issued underneath.

Every field named _ns is nanoseconds, and the two timestamps are nanoseconds since the Unix epoch, so comparable with time.time_ns(). They are measured on a monotonic clock and mapped through an anchor the monitor took when it was constructed, so a stepped system clock cannot corrupt any duration, while a long run may drift from the wall clock by whatever NTP did to it.

class BackendTotals[source]#

What one backend carried, a value of Summary.by_backend

num_ops: int#
bytes_transferred: int#
total_duration_ns: int#
num_errors: int#
start_unix_ns: int#

When counting started, or was last reset

end_unix_ns: int#

When the summary was read

num_ops: int#

Number of user-facing operations

num_reads: int#

Number of operations that were reads

num_writes: int#

Number of operations that were writes

bytes_requested: int#

Bytes the operations asked for

bytes_transferred: int#

Bytes actually transferred

Differs from bytes_requested on a short or failed read.

bytes_read: int#

Of the transferred bytes, how many were read

bytes_written: int#

Of the transferred bytes, how many were written

num_errors: int#

Number of operations that failed

busy_ns: int#

Time during which at least one operation was in flight

An approximation of the union of the operations’ spans: overlapping work is counted once, the gaps between calls are counted as idle, and it never exceeds wall_ns. An idle gap can be counted as busy when a finish reaches the monitor after a start that followed it, which takes two threads and a gap shorter than the delay between stamping a report and delivering it.

total_duration_ns: int#

The operations’ durations added up

Unlike busy_ns, which counts a stretch of time once however many operations filled it, this counts every operation. Only completed operations contribute.

by_backend: dict[str, BackendTotals]#

What each backend carried, keyed by the backend’s name

The totals partition the summary’s own, every operation belonging to exactly one backend. There is no per-backend busy time, that being a union over wall time which two backends running at once would both claim.

Excluded from hash(), as an unhashable field, and only from that. It still takes part in ==, so two summaries that differ here are unequal, they merely share a hash bucket.

counters: dict[str, int]#

The work in the span that belongs to no single operation

The counters run for the life of the process, and this is the part of them that falls inside the span.

Excluded from hash() for the same reason as by_backend.

wall_ns: int#

Wall-clock span this summary covers

Between start_unix_ns and end_unix_ns.

busy_bytes_per_sec: float#

Throughput while KvikIO was actually busy, or zero if no time was spent busy

Dividing the bytes by wall_ns instead would make a program that reads for 10 ms and then computes for 90 ms look ten times slower than its storage really is. Multiply by busy_fraction to recover the whole-span rate.

Understates while an operation is in flight, since its time counts from the moment it starts and its bytes only once it completes.

busy_fraction: float#

Fraction of the span during which KvikIO was doing something

Between 0 and 1. At 0.9 the program is nearly always doing I/O, at 0.03 it was idle almost throughout.

mean_duration_ns: int#

Average time one operation took, or zero if nothing completed

__init__(
start_unix_ns: int,
end_unix_ns: int,
num_ops: int,
num_reads: int,
num_writes: int,
bytes_requested: int,
bytes_transferred: int,
bytes_read: int,
bytes_written: int,
num_errors: int,
busy_ns: int,
total_duration_ns: int,
by_backend: dict[str, BackendTotals],
counters: dict[str, int],
wall_ns: int,
busy_bytes_per_sec: float,
busy_fraction: float,
mean_duration_ns: int,
) None#
since(
previous: Summary,
) Summary[source]#

Totals for the interval between an earlier reading and this one

Reporting periodically wants one reading per tick, differenced against the last. Two calls to SummaryMonitor.since() would leave a gap between them, and an operation that completed in the gap would fall into both intervals:

baseline = monitor.get()
while running:
    time.sleep(interval)
    now = monitor.get()
    report(now.since(baseline))
    baseline = now
Parameters:
previous

An earlier reading of the same span.

Returns:
The interval’s totals.
Raises:
ValueError

If previous is not an earlier reading of the same span, which covers an interval, a reading from another monitor, and one from before a reset.

to_json() str[source]#

Serialize to JSON

The timestamps are against the wall clock, so another program can line the summary up with its own log.

Returns:
A JSON object.
serialize() bytes[source]#

Serialize to bytes, exactly

Everything survives, including the clock anchor, so a summary that has been through a pipe is still a valid previous for SummaryMonitor.since(). Pickling uses this.

Returns:
A fixed-size buffer, which only this version of KvikIO reads back.
static deserialize(data: bytes) Summary[source]#

Rebuild a summary from serialize()

Parameters:
data

What serialize() produced.

Returns:
Summary

The summary.

Raises:
ValueError

If the bytes are not a summary, are truncated, or carry a version this build does not know.

report(all_rows: bool = False) str[source]#

Format a human-readable report of every field

Byte counts, durations and rates are scaled to readable units. Use to_json() instead when the output is going to be parsed.

Parameters:
all_rows

Print every row, including the backends the run never reached and the subsystems it never touched.

Returns:
The report, one field per line, newline-terminated.

Defaults#

class kvikio.defaults.ConfigContextManager(config: dict[str, str])[source]#

Context manager allowing the KvikIO configurations to be set upon entering a with block, and automatically reset upon leaving the block.

__init__(config: dict[str, str])[source]#
kvikio.defaults.set(
config: dict[str, Any],
/,
) ConfigContextManager[source]#
kvikio.defaults.set(
key: str,
value: Any,
/,
) ConfigContextManager

Set KvikIO configurations.

Examples:

  • To set one or more properties

    # Set the property globally.
    kvikio.defaults.set({"prop1": value1, "prop2": value2})
    
    # Set the property with a context manager.
    # The property automatically reverts to its old value
    # after leaving the `with` block.
    with kvikio.defaults.set({"prop1": value1, "prop2": value2}):
        ...
    
  • To set a single property

    # Set the property globally.
    kvikio.defaults.set("prop", value)
    
    # Set the property with a context manager.
    # The property automatically reverts to its old value
    # after leaving the `with` block.
    with kvikio.defaults.set("prop", value):
        ...
    
Parameters:
config

The configurations. Can either be a single parameter (dict) consisting of one or more properties, or two parameters key (string) and value (Any) indicating a single property.

Valid configuration names are:

  • "compat_mode"

  • "num_threads"

  • "task_size"

  • "gds_threshold"

  • "bounce_buffer_size"

  • "http_max_attempts"

  • "http_status_codes"

  • "http_timeout"

  • "auto_direct_io_read"

  • "auto_direct_io_write"

  • "remote_io_backend"

Returns:
ConfigContextManager

A context manager. If used in a with statement, the configuration will revert to its old value upon leaving the block.

kvikio.defaults.get(config_name: str) Any[source]#

Get KvikIO configurations.

Parameters:
config_name: str

The name of the configuration.

Valid configuration names are:

  • "compat_mode"

  • "num_threads"

  • "task_size"

  • "gds_threshold"

  • "bounce_buffer_size"

  • "http_max_attempts"

  • "http_status_codes"

  • "http_timeout"

  • "auto_direct_io_read"

  • "auto_direct_io_write"

  • "remote_io_backend"

Returns:
Any

The value of the configuration.

Nsight Systems plugin#

kvikio.nsys.nsys_plugin_search_dir() str[source]#

Return the directory that holds KvikIO’s bundled Nsight Systems plugins.

Works for both installation channels, which lay out the plugin differently.

  • pip (inside the libkvikio wheel, a Python package):

site-packages/libkvikio/
|-- __init__.py, load.py, nsys.py
|-- lib/libkvikio.so
|-- share/kvikio/nsys-plugins/ --> returned
    |-- kvikio_nic/{kvikio_nic_nsys_plugin, nsys-plugin.yaml}
  • conda (the libkvikio conda package, C++ files only, no Python module):

$CONDA_PREFIX/
|-- lib/libkvikio.so
|-- include/kvikio/
|-- share/kvikio/nsys-plugins/ --> returned
    |-- kvikio_nic/{kvikio_nic_nsys_plugin, nsys-plugin.yaml}

See profiling for how to enable the plugin in Nsight Systems.

Returns:
str

The plugin search directory. The path is returned even if it does not exist, for example when the plugin was not built.