注意

RAFT 中的向量搜索和聚类算法正在迁移到一个专门用于向量搜索的新库，名为 cuVS。在迁移过程中，我们将继续支持 RAFT 中的向量搜索算法，但在 RAPIDS 24.06（六月）版本发布后将不再更新。我们计划在 RAPIDS 24.10（十月）版本之前完成迁移，并在 24.12（十二月）版本中将其从 RAFT 中完全移除。

通用

本页提供 pylibraft 类参考，用于说明 pylibraft.common 包中公开的元素。

基本术语

class pylibraft.common.DeviceResources

DeviceResources 是一个轻量级的 Python 包装器，封装了 RAFT C++ 接口中公开的相应 C++ 类 device_resources。有关此结构的接口级别详细信息，请参阅头文件 raft/core/device_resources.hpp

参数:

stream可选流，用于排序 CUDA 指令

接受 pylibraft.common.Stream() 或 uintptr_t (cudaStream_t)

方法

getHandle(self)	返回底层 raft::device_resources 实例的指针，类型为 size_t
sync(self)	对此实例设置的流发出同步指令。

示例

基本用法

>>> from pylibraft.common import Stream, DeviceResources
>>> stream = Stream()
>>> handle = DeviceResources(stream)
>>>
>>> # call algos here
>>>
>>> # final sync of all work launched in the stream of this handle
>>> # this is same as `raft.cuda.Stream.sync()` call, but safer in case
>>> # the default stream inside the `device_resources` is being used
>>> handle.sync()
>>> del handle  # optional!

将 cuPy 流与 RAFT device_resources 一起使用

>>> import cupy
>>> from pylibraft.common import Stream, DeviceResources
>>>
>>> cupy_stream = cupy.cuda.Stream()
>>> handle = DeviceResources(stream=cupy_stream.ptr)

将 RAFT 流与 CuPy ExternalStream 一起使用

>>> import cupy
>>> from pylibraft.common import Stream
>>>
>>> raft_stream = Stream()
>>> cupy_stream = cupy.cuda.ExternalStream(raft_stream.get_ptr())

getHandle(self)

返回底层 raft::device_resources 实例的指针，类型为 size_t

sync(self)

对此实例设置的流发出同步指令。

class pylibraft.common.Stream

Stream 是 cudaStream_t 及其操作的轻量级包装器。

方法

get_ptr(self)	返回底层 cudaStream_t 句柄的 uintptr_t 指针
sync(self)	在此对象拥有的 cuda 流上同步。

示例

>>> from pylibraft.common.cuda import Stream
>>> stream = Stream()
>>> stream.sync()
>>> del stream  # optional!

get_ptr(self)

返回底层 cudaStream_t 句柄的 uintptr_t 指针

sync(self)

在此对象拥有的 cuda 流上同步。请注意，由于先前异步启动的问题，这可能会引发异常

class pylibraft.common.device_ndarray(np_ndarray)

pylibraft.common.device_ndarray 是一个非常轻量级的 __cuda_array_interface__ 包装器，封装了 numpy.ndarray。

属性:

c_contiguous

当前的 device_ndarray 是否按行主序排列？

dtype

当前 device_ndarray 实例的数据类型

f_contiguous

当前的 device_ndarray 是否按列主序排列？

shape

当前 device_ndarray 实例的形状

strides

当前 device_ndarray 实例的步幅

方法

property c_contiguous

当前的 device_ndarray 是否按行主序排列？

copy_to_host()

返回一个在主机上新的 numpy.ndarray 对象，其中包含此 device_ndarray 的当前内容

property dtype

当前 device_ndarray 实例的数据类型

classmethod empty(shape, dtype=<class 'numpy.float32'>, order='C')

返回一个具有给定形状和类型的新 device_ndarray，不初始化条目。

参数:

shapeint 或 int 元组

空数组的形状，例如 (2, 3) 或 2。

dtype数据类型，可选

数组所需的输出数据类型，例如 numpy.int8。默认值为 numpy.float32。

order{‘C’，‘F’}，可选（默认值：‘C’）

是否以行主序（C 风格）或列主序（Fortran 风格）在内存中存储多维数据

property f_contiguous

当前的 device_ndarray 是否按列主序排列？

property shape

当前 device_ndarray 实例的形状

property strides

当前 device_ndarray 实例的步幅

可中断

pylibraft.common.interruptible.cuda_interruptible()

临时安装一个键盘中断处理程序 (Ctrl+C)，该处理程序会取消封闭的可中断 C++ 线程。

在通过 cython 导入的长时间运行的 C++ 函数上使用此功能

>>> with cuda_interruptible():
>>>     my_long_running_function(...)

还建议在调用期间释放 GIL，以确保处理程序有机会运行

>>> with cuda_interruptible():
>>>     with nogil:
>>>         my_long_running_function(...)

pylibraft.common.interruptible.synchronize(Stream stream: Stream)

与 cudaStreamSynchronize 相同，但在 with cuda_interruptible() 块内调用时可被中断。

pylibraft.common.interruptible.cuda_yield()

检查异步接收到的 interrupted_exception。如果在之前用户在 with cuda_interruptible() 块内按下 Ctrl+C，则引发该异常。

CUDA 数组接口辅助函数

class pylibraft.common.cai_wrapper(cai_arr)

围绕 CUDA 数组接口对象的简单包装器，用于减少从底层字典中提取常见信息的模板代码。