Skip to content

Python API should handle commonly used CUDA stream objects #895

@kingcrimsontianyu

Description

@kingcrimsontianyu

Some KvikIO Python API functions have CUDA stream parameters.

def raw_read_async(
        self,
        buf,
        raw_stream: int,
        size: Optional[int] = None,
        file_offset: int = 0,
        dev_offset: int = 0,
) -> IOFutureStream:

def stream_register(raw_stream: int, flags: int) -> None:

Currently, the stream argument passed to these functions has to be raw handles of integer type (usually the underlying CUDA stream, which is a pointer). To make the API more flexible and robust, we may consider supporting common types of CUDA stream objects, on top of the existing raw CUDA stream support:

It is likely that we do not need to include these packages as KvikIO dependencies. Instead we only need to require certain methods to be available with which to retrieve the underlying CUDA stream and pass it to Cython.

Essentially KvikIO is expected to have a stream wrapper, say kvikio.Stream(obj=external_stream_obj_or_raw_stream) that retrieves the underlying CUDA stream from the abovementioned stream objects.

Metadata

Metadata

Assignees

No one assigned

    Labels

    feature requestNew feature or requestpythonAffects the Python API of KvikIO

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions