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 ベースのローダー、形式固有のデコーダー、サイドカーメタデータ、部分読み込み / シャーデッド読み込み）が必要でない限り、組み込みローダーを優先してください。

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

.npy はヘッダーに shape と dtype がディスク上に保存されるため、具体例として取り上げています。ただし、このひな形は、範囲 / スライス読み込みのコストが低い任意の形式（raw binary、HDF5、Parquet/Arrow — 下記の「その他の形式」を参照）に適用できます。リファレンス実装： assets/examples/parallel_npy_load.py。

データレイアウトの前提

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

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

• shard_0000.npy、shard_0001.npy、... のように、連続した整数のファイル名（幅 4 のゼロ埋め）を持つファイルが格納されたディレクトリです。
• すべてのシャードで dtype と後続の軸（shape[1:]）が一致します。axis 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,))