cupynumeric-parallel-data-load

並列シャーディングデータから cupynumeric への読み込み

このスキルがある理由。 cupynumeric は単一の .npy ファイルに対する cupynumeric.load を含め、NumPy の配列 API を踏襲しています。それ以外のファイルの読み込みは cupynumeric ではなく Legate が担います。

形式	組み込みローダー
単一の .npy	cupynumeric.load(path)（NumPy API との互換性）
HDF5（単一ファイル）	legate.io.hdf5.from_file / from_file_batched
シャーディングされた複数ファイル（任意の形式）、Parquet/Arrow、raw binary、カスタムレイアウト	組み込みローダーなし — このスキルを使います。

このスキルでは、最後の行にある不足を補う標準的な方法を紹介します。必要な形式用のサードパーティ製リーダー（h5py、pyarrow、np.memmap など）をタスク本体の中で呼び出す Legate Python タスクを作成し、Legate に読み込みを GPU／ノード間で分散させます。組み込みローダーがある形式では、タスク内のカスタム処理（mmap ベースのローダー、形式固有のデコーダー、サイドカーメタデータ、部分読み込み／シャーディング読み込み）が必要な場合を除き、そちらを優先してください。

標準パターンは、ファイルではなくマシンに合わせた手動パーティション分割と手動タスク起動です。シャーディングするのは軸 0 のみで、後続の軸は各タイルに含めます。ファイルごとにシャードの行数が異なっていても構いません（dtype と後続の軸は一致している必要があります）。起動時には、ファイル数にかかわらず、利用可能なすべてのプロセッサーを使います。

.npy を実例に使うのは、ヘッダーに形状と dtype が記録されているためです。ただし、この骨組みは、範囲／スライスの読み込みが低コストな任意の形式（raw binary、HDF5、Parquet/Arrow。「その他の形式」を参照）に適用できます。リファレンス実装： assets/examples/parallel_npy_load.py。

データレイアウトの前提

このスキルは読み込みのみを扱います。データが、予測可能でインデックス指定可能な形式で共有ファイルシステム上にすでに配置されていることが前提です。ファイルの生成は対象外です（利便性のため、サンプルには write サブコマンドが含まれていますが、実際のユーザーは独自のファイルを用意します）。

実例では、次の特定のレイアウトを前提としています。

• shard_0000.npy、shard_0001.npy、…の名前のファイルが、連続した整数の並び（4 桁のゼロ埋め）で格納されたディレクトリ。
• すべてのシャードで dtype と後続の軸（shape[1:]）が一致していること。軸 0（シャードあたりの行数）はファイル間で異なっていても構いません。この手順では累積行オフセットのテーブルを作成し、リーフタスク内で各ファイルの重複するスライスを読み込みます。
• すべてのランクからディレクトリが参照できること（マルチノード実行では共有ファイルシステムを使用）。

サンプルの discover_layout() は検出した内容を出力し、レイアウトが誤っている場合（ディレクトリがない、シャードがない、dtype／後続の軸が一致しない、連続した shard_NNNN.npy の並びに抜けがある）には、説明的なエラーを出して処理を停止します。

データのレイアウトが異なる場合（固定ストライドの raw binary、シャードごとに 1 つのデータセットを含む 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 を導出します。

--min-gpu-chunk 1 が必要なのは、タイルごとの要素数が GPU 起動時の Legate のデフォルト最小チャンクサイズを下回る場合のみです（例：動作例のデフォルト設定では、合計行数を 4 基の GPU に分割したときのタイルあたりの要素数が ~1M となり、しきい値を下回るため、指定しないと 1 基の GPU にまとめられます）。本番規模のデータセット（タイルあたり数千万以上の要素）では、このフラグを省略して Legate のデフォルト設定を使用できます。各タイルが十分に大きく、すべての GPU にタイルを割り当てることよりもタスクごとのオーバーヘッドが重要な場合は、中程度の値（例：--min-gpu-chunk 1024）に引き上げても問題ありません。

手順

.npy の動作例を使った 5 つの手順です。形式固有なのは、手順 1（形式ヘッダーの解析）と手順 4（タスク本体内のファイルごとのリーダー）のみです。残りの 3 つ（出力先の確保、分割、同期）は形式を問わずそのまま再利用できます。形式の切り替え箇所については、以下の「その他の形式」を参照してください。

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 は、起動前にドライバーによってモジュールグローバル変数として設定されます。制御レプリケーションではドライバーがすべてのランクで実行されるため、すべてのワーカーが同一の値を参照できます。

各タスクはまずコンシューマービュー（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:]）を一致させる必要があります。 このレシピではシャードを axis 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 は共有ファイルシステム上にある必要があります。 すべてのワーカー（すべてのランク）がパスを使ってシャードを開きます。ノードローカルの /tmp パスが使えるのは、単一ノードのデモに限られます。

バリアント

均一シャードの高速パス（ファイルごとに 1 タスク）

すべてのシャードがすでに同じ (shape, dtype) を持ち、たまたま num_shards 個のプロセッサーを利用できる場合は、cum-rows / 二分探索の処理がオーバーヘッドになります。tile_rows = shard_shape[0]、num_tasks = num_shards を設定すると、パーティションにはファイルごとに 1 つのタイルが割り当てられ、各タスクが 1 つのファイルを最初から最後まで読み込みます（二分探索も内側のループもありません）。ドライバー側の切り替えは 1 行で済みます。

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 タスク本体はどちらのモードでも動作します。内側のループがタスクごとにちょうど 1 回だけ反復するだけです。高速パス専用のタスク本体を別途用意する必要はありません。

過剰分割による負荷分散の改善

デフォルトの tile_rows = ceil(total_rows / num_processors) では、プロセッサーごとに 1 つのタイルが割り当てられます。係数 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) を受け取り、C 連続にできる形状 (file_hi - file_lo,) + trailing_shape の numpy 配列を返すことです。 低コストで範囲／スライスを読み取れる必要があります。「ファイル全体を 読み取る」形式しかサポートしない場合、部分的な重なりのケース（1 つの ファイルの一部だけをカバーするタイル）に対応できません。

形式	リーフタスク内のリーダー
.npy（動作例）	host = np.ascontiguousarray(np.load(p, mmap_mode="r")[file_lo:file_hi])
Raw binary（固定形状）	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

（各形式の組み込み単一呼び出しローダーについては、このファイルの先頭にある「このスキルが存在する理由」表を参照してください。）

検出ステップ（ステップ 1）では、各形式のメタデータを解析します。.npy / HDF5 / Parquet は、ファイルごとの行数と dtype をディスク上に保持しています。 Raw binary には保持されていないため、サイドカーファイルを使うか、ファイルサイズから導出します。

よくある落とし穴

リーフタスク内での 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 カプセルを処理します。cn.asarray は ドライバー内では問題ありませんが、リーフタスク内では使えません。torch を使った回避策については、examples/dlpack/leaf_task_interop.py を参照してください。

タスク内の 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,))