cupynumeric-parallel-data-load

并行分片数据 → cupynumeric 加载

为什么需要此技能。 cupynumeric 镜像了 NumPy 的数组 API，包括用于加载单个 .npy 文件的 cupynumeric.load。除此之外，文件加载功能由 Legate 提供，而非 cupynumeric：

格式	内置加载器
单个 .npy	cupynumeric.load(path)（与 NumPy API 保持一致）
HDF5（单个文件）	legate.io.hdf5.from_file / from_file_batched
分片多文件（任意格式）、Parquet/Arrow、原始二进制、自定义布局	没有内置加载器——请使用此技能。

此技能介绍如何以规范方式弥补表格最后一行所述的不足：编写一个 Legate Python 任务，在任务体内调用相应格式所需的第三方读取器（h5py、pyarrow、np.memmap 等），并让 Legate 将读取任务分配到各个 GPU / 节点。对于已有内置加载器的格式，除非你需要自定义任务体（基于 mmap 的加载器、特定格式的解码器、辅助元数据、部分 / 分片读取），否则请优先使用内置加载器。

规范模式：手动分区 + 手动启动任务，按机器配置而非文件数量确定规模。 只有轴 0 会分片；其余轴会随每个数据块一起处理。不同文件的每个分片行数可以不同（只需 dtype 和其余轴一致）；无论文件数量多少，启动任务都会用满所有可用处理器。

.npy 是本文的示例格式，因为其形状和 dtype 信息存储在文件头中，但该框架适用于任何支持低开销范围 / 切片读取的格式（原始二进制、HDF5、Parquet/Arrow——请参阅下文的“其他格式”）。参考实现： assets/examples/parallel_npy_load.py。

数据布局假设

此技能仅涉及加载——假设数据已按某种可预测、可索引的方式存放在共享文件系统中。如何生成这些文件不在讨论范围内（示例为方便起见提供了 write 子命令，但实际用户可使用自己的数据）。

示例假设采用一种特定布局：

• 一个目录，其中包含名为 shard_0000.npy、shard_0001.npy、……的文件，编号为连续的整数（宽度为 4，不足位数以零填充）。
• 所有分片的 dtype 和其余轴（shape[1:]）都相同；轴 0（每个分片的行数）可以不同——此方案会构建累积行偏移表，并在叶任务中读取每个文件与目标区域重叠的切片。
• 每个 rank 都可以访问该目录（多节点运行时使用共享文件系统）。

示例中的 discover_layout() 会打印检测结果；若布局不正确（目录不存在、没有分片、dtype / 其余轴不匹配，或连续的 shard_NNNN.npy 序列中存在缺号），则会以明确的错误信息终止运行。

如果你的数据采用其他布局——固定步长原始二进制、每个分片对应一个数据集的 HDF5 文件、目录树等——只需修改 glob 模式、每个文件的读取器（下文步骤 4）和元数据发现逻辑（下文步骤 1）。分区与启动任务的机制与布局无关。

适用场景

格式选择请参阅上方的格式表（使用内置加载器还是此技能）。此外，以下两种情况也适合使用此技能：

• 将顺序执行的 np.concatenate([read(f) for f in files]) 替换为按 GPU 并行读取。
• 演示如何通过手动启动任务，让用户自定义的 Legate Python 任务写入 cupynumeric 输出数组。

示例

下面的路径均相对于此技能目录（脚本位于 assets/examples/parallel_npy_load.py）。请调整前缀，使其与技能的安装位置相符（例如，如果技能位于顶层 skills/ 目录下，则使用 skills/cupynumeric-parallel-data-load/assets/...）。

# Single-node, 4 GPUs.
legate --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

# Multi-node, 2 nodes x 4 GPUs (slurm), shared filesystem at --shard-dir.
# Generate the shards once on rank 0, then re-run `read` at any scale.
legate --launcher srun --nodes 2 --cpus 1 \
    assets/examples/parallel_npy_load.py \
    write --shard-dir /shared/scratch/demo

legate --launcher srun --nodes 2 --ranks-per-node 4 \
    --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

不设置布局标志时，读取驱动程序会遍历每个 .npy 文件头，以获取每个文件的行数、尾部形状和 dtype，然后根据可用处理器数量计算 tile_rows。

只有当每个分块的元素数低于 Legate 对 GPU 启动设置的默认最小分块大小时，才需要设置 --min-gpu-chunk 1（例如，演示示例中使用默认设置时，将总行数分配到 4 个 GPU，每个分块约有 1M 个元素，低于阈值；否则这些数据会被合并到单个 GPU 上）。对于生产规模的数据集（每个分块包含数千万个或更多元素），可以省略此标志，让 Legate 使用默认值。如果每个分块足够大，以至于每个任务的开销比让每个 GPU 都获得一个分块更重要，也可以将其调高到适中的值（例如 --min-gpu-chunk 1024）。

说明

以下是基于 .npy 示例的五个步骤；只有第 1 步（解析格式文件头）和第 4 步（任务体内逐文件读取器）与格式有关。其他三步（分配目标、分区、栅栏）可在不同格式间原样复用——有关需要替换的部分，请参见下文“其他格式”。

1. 读取每个分片的元数据

扫描目录并查看每个 .npy 文件头（mmap_mode="r" 只读取文件头）。文件头包含每个分片的形状和 dtype，因此驱动程序无需加载数据，即可获取总行数、尾部形状和累积行偏移表：

paths = sorted(SHARD_DIR.glob("shard_*.npy"))

per_file_rows = []                       # rows along axis 0 per file
trailing_shape = None                    # shape[1:], must match across files
dtype = None
for p in paths:
    hdr = np.load(p, mmap_mode="r")
    if trailing_shape is None:
        trailing_shape = tuple(hdr.shape[1:])
        dtype = hdr.dtype
    elif tuple(hdr.shape[1:]) != trailing_shape or hdr.dtype != dtype:
        raise RuntimeError(
            f"{p.name}: trailing shape / dtype mismatch "
            f"({hdr.shape[1:]}/{hdr.dtype} vs {trailing_shape}/{dtype})"
        )
    per_file_rows.append(int(hdr.shape[0]))

cum_rows = np.cumsum([0] + per_file_rows, dtype=np.int64)  # length N+1
total_rows = int(cum_rows[-1])

上面的代码片段会检查各文件的 dtype 和 trailing_shape（即 shape[1:]）是否一致。每个分片的行数可以不同——累积行数表会处理这种情况。生产代码还应验证文件名是否构成连续序列 shard_0000.npy ... shard_NNNN.npy（为简洁起见，代码片段中省略了此步骤；请参见示例中的 discover_layout()）。数据发现仅依赖磁盘格式本身公开的信息（此处为 .npy 文件头；HDF5 则为 .shape / .dtype 等）；任何旁路文件（清单、内容哈希）都属于额外的验证步骤。

2. 根据元数据创建 cupynumeric 输出存储

数组沿轴 0 的总长度为 total_rows；尾部各轴的形状直接采用 trailing_shape。使用 cn.empty——任务会覆盖每个单元格，初始化为零纯属浪费。

import cupynumeric as cn

total_shape = (total_rows,) + trailing_shape
out = cn.empty(total_shape, dtype=dtype)

3. 按处理器数量对存储进行分块

启动形状根据可用处理器数量确定，而不是根据文件数量。选择 tile_rows = ceil(total_rows / num_processors)，并按该分块大小对轴 0 进行分区。尾部各轴不进行分区（分块会覆盖这些轴的全部范围）。最后一个分块可以较短——这正是 partition_by_tiling 支持的情况——因此此方案不要求长度可整除。

from legate.core import TaskTarget, get_legate_runtime
from legate.core.data_interface import as_logical_array

runtime = get_legate_runtime()
machine = runtime.get_machine()
num_processors = max(
    machine.count(TaskTarget.GPU),
    machine.count(TaskTarget.OMP),
    machine.count(TaskTarget.CPU),
    1,
)

tile_rows = max(1, (total_rows + num_processors - 1) // num_processors)
tile_shape = (tile_rows,) + trailing_shape
partition = as_logical_array(out).data.partition_by_tiling(tile_shape)

num_tasks = (total_rows + tile_rows - 1) // tile_rows  # match partition tile count

4. 定义叶子任务并手动启动

驱动程序在启动前会将 PATHS、CUM_ROWS（第 1 步中的文件路径和累计行偏移表）以及 TILE_ROWS 填充为模块全局变量；控制复制会在每个 rank 上运行驱动程序，因此每个工作进程看到的值都相同。

每个任务会先构建其使用方视图（GPU 上使用 cupy，CPU/OMP 上使用 numpy），并从 view.shape[0] 读取该分块的实际行数——PhysicalStore 本身没有 .shape 属性，因此必须通过视图访问。然后，任务根据其启动坐标和行数计算全局行范围，对 cum_rows 执行二分查找，找出与之重叠的文件，并将每个重叠文件的切片复制到对应的目标切片中。注册 CPU、OMP 和 GPU 变体，这样同一启动在任何地方都无需修改即可运行；根据 ctx.get_variant_kind() 分派，选择与 OutputStore 所在位置匹配的使用方（FBMEM 使用 cp.from_dlpack(dst)，SYSMEM 使用 np.asarray(dst)）。cupy 仅在 GPU 分支内导入，因此在没有 cupy 的机器上也能加载任务主体。

import bisect
from legate.core import TaskContext, VariantCode
from legate.core.task import OutputStore, task

@task(variants=(VariantCode.CPU, VariantCode.OMP, VariantCode.GPU))
def load_tile(ctx: TaskContext, dst: OutputStore) -> None:
    t = ctx.task_index[0]                              # tile index 0..num_tasks-1

    variant = ctx.get_variant_kind()
    if variant == VariantCode.GPU:
        import cupy as cp                              # lazy: only on GPU
        view = cp.from_dlpack(dst)
    else:
        view = np.asarray(dst)                         # zero-copy numpy view

    tile_rows_actual = view.shape[0]                   # short on the last tile
    row_start = t * TILE_ROWS                          # global axis-0 start
    row_end = row_start + tile_rows_actual

    # Find the half-open range of file indices that overlap [row_start, row_end).
    first_file = bisect.bisect_right(CUM_ROWS, row_start) - 1
    last_file = bisect.bisect_right(CUM_ROWS, row_end - 1) - 1

    for f in range(first_file, last_file + 1):
        # Intersection of tile [row_start, row_end) with file [cum[f], cum[f+1]).
        lo = max(row_start, int(CUM_ROWS[f]))
        hi = min(row_end, int(CUM_ROWS[f + 1]))
        file_lo = lo - int(CUM_ROWS[f])
        file_hi = hi - int(CUM_ROWS[f])
        dst_lo = lo - row_start
        dst_hi = hi - row_start
        chunk = np.ascontiguousarray(
            np.load(PATHS[f], mmap_mode="r")[file_lo:file_hi]
        )
        if variant == VariantCode.GPU:
            view[dst_lo:dst_hi].set(chunk)             # cudaMemcpyAsync H2D
        else:
            view[dst_lo:dst_hi] = chunk                # zero-copy numpy write

manual_task = runtime.create_manual_task(
    load_tile.library,
    load_tile.task_id,
    (num_tasks,),                                      # launch domain == tile count
)
manual_task.add_output(partition)
manual_task.execute()

两种使用方都会通过 PhysicalStore 的原生生产者访问数据（cupy 使用 __dlpack__，np.asarray 使用 __array_interface__），生成指向本地分块的零拷贝视图。二分查找的开销为 O(log num_shards)，内部循环通常迭代 1–2 次（分块最多与少数几个文件重叠）。

5. 同步并验证

get_legate_runtime().issue_execution_fence(block=True)

严格约束

1. 所有分片的 dtype 和尾随轴（shape[1:]）必须相同。 此方案沿轴 0 堆叠分片；目标的尾随轴来自 trailing_shape，该值在发现步骤中锁定为第一个文件的值。每个分片的行数（shape[0]）可以不同，累计偏移表会处理这一情况。示例会在任何分片的 dtype 或尾随形状与第一个分片不同时，报出描述性错误。

2. 选择与变体匹配的使用方。 cp.from_dlpack 不接受驻留在 SYSMEM 中的存储；np.asarray 则会静默返回驻留在 FBMEM 中的存储的主机视图，而你实际上无法通过该视图写入。应根据 ctx.get_variant_kind() 分派，让每个变体使用自己的使用方——参见第 4 步。

3. mmap 视图不一定是 C 连续的——在调用 .set() 或执行 numpy 原地写入之前，使用 np.ascontiguousarray(arr[file_lo:file_hi]) 包装每个文件切片。

4. 多节点：SHARD_DIR 必须位于共享文件系统上。 每个工作进程（每个 rank 上的工作进程）都会按路径打开分片；节点本地的 /tmp 路径仅适用于单节点演示。

变体

均匀分片快速路径（每个文件一个任务）

如果每个分片的 (shape, dtype) 已相同，且恰好有 num_shards 个可用处理器，那么累计行数和二分查找机制就会带来额外开销。将 tile_rows = shard_shape[0]、num_tasks = num_shards；这样分区就会为每个文件分配一个分块，每个任务恰好从头到尾读取一个文件（无需二分查找，也无需内部循环）。驱动程序侧只需一行即可切换：

if all(r == per_file_rows[0] for r in per_file_rows) and num_shards == num_processors:
    tile_rows = per_file_rows[0]
else:
    tile_rows = max(1, (total_rows + num_processors - 1) // num_processors)

无论采用哪种模式，同一个 load_tile 任务主体都可以运行——内部循环只是在每个任务中恰好迭代一次。快速路径无需单独的任务主体。

过度分解以改善负载均衡

默认的 tile_rows = ceil(total_rows / num_processors) 会为每个处理器分配一个分块。若要按 K 倍进行过度分解（分块更小、点任务更多、队列粒度更细），将分母改为 K * num_processors：

tile_rows = max(1, (total_rows + K * num_processors - 1) // (K * num_processors))

这样，num_tasks = ceil(total_rows / tile_rows) 就会扩展到大约 K * num_processors。同一个任务主体仍然适用——二分查找只会将更多任务定位到每个文件上。

其他格式

只需更改 load_tile 中的逐文件读取器。读取器的约定是：给定文件路径和沿轴 0 的半开行范围 [file_lo, file_hi)，返回一个形状为 (file_hi - file_lo,) + trailing_shape 的 numpy 数组，并且该数组可以转换为 C 连续布局。 必须支持低开销的范围／切片读取——只支持“读取整个文件”的格式无法处理部分重叠场景（即一个分块只覆盖某个文件的一部分）。

格式	叶任务中的读取器
.npy（示例）	host = np.ascontiguousarray(np.load(p, mmap_mode="r")[file_lo:file_hi])
原始二进制（固定形状）	arr = np.memmap(p, dtype=DTYPE, mode="r", shape=(rows_in_file, *trailing_shape)); host = np.ascontiguousarray(arr[file_lo:file_hi])
HDF5	with h5py.File(p, "r") as f: host = np.ascontiguousarray(f["data"][file_lo:file_hi])
Parquet / Arrow	tbl = pq.read_table(p, columns=..., use_threads=False).slice(file_lo, file_hi - file_lo); host = tbl.to_pandas().values

（各格式内置的单次调用加载器，请参阅本文件顶部的“Why this skill exists”表格。）

发现步骤（步骤 1）会解析各格式的元数据：.npy / HDF5 / Parquet 都会在磁盘上保存每个文件的行数和 dtype。 原始二进制则不会——需要使用边车文件或根据文件大小推导。

常见问题

在叶任务中使用 cn.asarray(dst) 是非法的

在 @task 函数体内，任何会触及顶层运行时的 cupynumeric 操作——cn.asarray(store)、切片赋值 cn_dst[s] = host_np——都会从错误的上下文触发 create_index_space，导致 Legion 中止：

LEGION API USAGE EXCEPTION: Invalid task context passed to runtime call
create_index_space

解决方法：在叶任务中使用第三方库（cupy / torch / numpy）消费 DLPack capsule。在驱动程序中可以使用 cn.asarray，但不能在叶任务中使用。参阅 examples/dlpack/leaf_task_interop.py，了解使用 torch 的变通方案。

任务内的 assert 会导致运行时中止

Legate 会将 @task 中未抛出的异常视为契约违规，并中止运行时，除非任务已通过 throws_exception() 注册。启动任务前，请在主机端进行健全性检查。

启动域必须与分区分块数量一致

create_manual_task(launch_shape=...) 和 partition_by_tiling(...) 相互独立——运行时不会检测二者不匹配。启动域较大 → 分块索引越界；启动域较小 → 部分分块未写入。始终基于同一组 (total_rows, tile_rows) 分别通过两次 ceil 运算推导二者（若直接将启动域大小设为 num_processors，当 num_processors > total_rows 时会启动过多任务）：

tile_rows = max(1, (total_rows + num_processors - 1) // num_processors)
num_tasks = (total_rows + tile_rows - 1) // tile_rows
partition = ...partition_by_tiling((tile_rows,) + trailing_shape)
runtime.create_manual_task(load_tile.library, load_tile.task_id, (num_tasks,))