Données fragmentées en parallèle -> chargement via cupynumeric

Pourquoi cette compétence existe-t-elle ? cupynumeric reproduit l'API des tableaux de NumPy, y compris la fonction cupynumeric.load pour un seul fichier .npy. Au-delà de cela, le chargement des fichiers s'effectue dans Legate, et non dans cupynumeric :

Cette compétence montre la méthode canonique pour combler la lacune de la dernière ligne : écrivez une tâche Python Legate qui appelle le lecteur tiers requis par le format (h5py, pyarrow, np.memmap, ...) à l'intérieur du corps de la tâche, et laissez Legate répartir les lectures entre les GPU / nœuds. Pour les formats disposant d’un chargeur intégré, privilégiez-le sauf si vous avez besoin d’un corps de tâche personnalisé (chargeur basé sur mmap, décodeur spécifique au format, métadonnées sidecar, lectures partielles / fragmentées).

Modèle canonique : partition manuelle + lancement manuel de la tâche, dimensionné en fonction de la machine, et non des fichiers. Seul l’axe 0 est fragmenté ; les axes suivants sont regroupés au sein de chaque tuile. Le nombre de lignes par fragment peut varier d’un fichier à l’autre (seuls le type de données et les axes suivants doivent correspondre) ; le lancement occupe tous les processeurs disponibles, quel que soit le nombre de fichiers.

Le format .npy sert d’exemple pratique car l’en-tête contient la forme et le type de données sur le disque, mais le principe s’applique à tout format permettant des lectures de plage/tranche peu coûteuses (binaire brut, HDF5, Parquet/Arrow — voir « Autres formats » ci-dessous). Implémentation de référence : assets/examples/parallel_npy_load.py.

Hypothèse concernant la disposition des données

Cette compétence concerne uniquement le chargement — elle suppose que les données sont déjà disposées sur un système de fichiers partagé de manière prévisible et indexable. La création de ces fichiers n’entre pas dans le champ d’application (l’exemple fournit une sous-commande d’écriture pour plus de commodité, mais les utilisateurs réels doivent apporter la leur).

L’exemple pratique suppose une structure spécifique :

• un répertoire contenant des fichiers nommés shard_0000.npy, shard_0001.npy, … selon une séquence d’entiers contiguë (largeur de 4 complétée par des zéros).
• Tous les shards partagent le même dtype et les mêmes axes de fin (shape[1:]) ; l’axe 0 (nombre de lignes par shard) peut varier d’un fichier à l’autre — la recette construit une table cumulative de décalages de lignes et lit la tranche chevauchante de chaque fichier depuis la tâche feuille.
• Le répertoire est accessible à tous les rangs (système de fichiers partagé pour les exécutions multi-nœuds).

La fonction `discover_layout()` de l’exemple affiche ce qu’elle a trouvé et génère une erreur fatale accompagnée d’un message descriptif lorsque la disposition est incorrecte (répertoire manquant, absence de fragments, type de données ou axes finaux non correspondants, ou une lacune dans la séquence contiguë `shard_NNNN.npy `).

Si vos données sont organisées selon une structure différente — format binaire brut à pas fixe, fichier HDF5 avec un ensemble de données par fragment, arborescence de répertoires, etc. — seuls le motif glob, le lecteur par fichier (étape 4 ci-dessous) et la découverte des métadonnées (étape 1 ci-dessous) changent. Le mécanisme de partitionnement et de lancement est indépendant de la structure des données.

Quand l’utiliser

Consultez le tableau des formats ci-dessus pour déterminer le choix du routage (chargeur intégré ou cette compétence). Au-delà de cela, voici deux autres indices indiquant que cette compétence est la solution adaptée :

• Remplacer la séquence np.concatenate([read(f) for f in files]) par des lectures parallèles par GPU.
• Démontrer comment une tâche Python Legate définie par l’utilisateur écrit dans un tableau de sortie cupynumeric via un lancement manuel.

Exemples

Les chemins d’accès ci-dessous sont écrits par rapport au répertoire de cette compétence (le script se trouve dans assets/examples/parallel_npy_load.py). Adaptez le préfixe en fonction de l’emplacement d’installation de votre skill (par exemple, skills/cupynumeric-parallel-data-load/assets/… si la skill se trouve dans un répertoire de niveau supérieur « skills/ » ).

# Un seul nœud, 4 GPU.
legate --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

# Multi-nœuds, 2 nœuds × 4 GPU (SLURM), système de fichiers partagé à l'emplacement --shard-dir.
# Générer les segments une seule fois sur le nœud 0, puis relancer `read` à n'importe quelle échelle.
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

Aucun indicateur de disposition — le pilote de lecture parcourt chaque en-tête .npy pour récupérer le nombre de lignes par fichier, la forme finale et le type de données, puis déduit le nombre de lignes par tuile (tile_rows) à partir du nombre de processeurs disponibles.

L’option --min-gpu-chunk 1 n’est nécessaire que lorsque le nombre d’éléments par tuile est inférieur à la taille minimale par défaut des chunks définie par Legate pour les lancements sur GPU (par exemple, les valeurs par défaut de l’exemple pratique — nombre total de lignes réparties sur 4 GPU à environ 1 million par tuile — tombent en dessous du seuil et seraient sinon regroupées sur un seul GPU). Pour les jeux de données de taille industrielle (des dizaines de millions d’éléments par tuile ou plus), vous pouvez omettre ce drapeau et laisser Legate utiliser sa valeur par défaut. Le régler sur une valeur modérée (par exemple --min-gpu-chunk 1024) est tout à fait acceptable lorsque chaque tuile est suffisamment grande pour que la surcharge par tâche prime sur l’attribution d’une tuile à chaque GPU.

Instructions

Cinq étapes tirées d’un exemple pratique au format .npy; seules l’étape 1 (analyse de l’ en-tête de format) et l’étape 4 (le lecteur par fichier à l’intérieur du corps de la tâche) sont spécifiques au format. Les trois autres (allocation de la destination, partition, barrière) sont réutilisées telles quelles quel que soit le format — voir « Autres formats » ci-dessous pour les points de basculement.

1. Lire les métadonnées de chaque fragment

Parcourez le répertoire et examinez chaque en-tête .npy (mmap_mode="r" ne lit que l’en-tête). L’en-tête contient la forme et le type de données par fragment, ce qui permet au pilote de récupérer le nombre total de lignes, la forme de fin et un tableau cumulatif de décalages de lignes sans jamais charger les données :

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

per_file_rows = []                       # nombre de lignes le long de l’axe 0 par fichier
trailing_shape = None                    # shape[1:], doit correspondre d’un fichier à l’autre
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} : incompatibilité entre la forme finale et le type de données "
            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)  # longueur N+1
total_rows = int(cum_rows[-1])

L'extrait de code ci-dessus garantit la correspondance entre le dtype et le trailing_shape (c'est-à-dire shape[1:]) d'un fichier à l'autre. Le nombre de lignes par shard peut varier — le tableau cum_rows gère cette situation. Le code de production doit également vérifier que les noms forment une séquence contiguë shard_0000.npy … shard_NNNN.npy (omise dans l’extrait par souci de concision ; voir discover_layout() dans l’ exemple détaillé). La détection repose uniquement sur ce que le format sur disque lui-même expose (l’en-tête .npy ici, .shape / .dtype pour HDF5, etc.) ; tout fichier d’accompagnement (manifeste, hachages de contenu) fait l’objet d’une étape de vérification distincte en plus.

2. Créer le magasin de sortie cupynumeric à partir des métadonnées

Le tableau total s’étend sur total_rows le long de l’axe 0 ; les axes suivants proviennent de trailing_shape sans modification. Utilisez cn.empty — la tâche écrase chaque cellule, une initialisation à zéro serait inutile.

import cupynumeric as cn

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

3. Diviser le magasin en tuiles en fonction du nombre de processeurs

La forme de lancement est dimensionnée en fonction du nombre de processeurs disponibles, et non du nombre de fichiers. Choisissez tile_rows = ceil(total_rows / num_processors) et partitionnez l’axe 0 selon cette taille de tuile. Les axes suivants ne sont pas partitionnés (la tuile s’étend sur toute leur longueur). La dernière tuile peut être plus courte — c’est exactement ce que prend en charge partition_by_tiling —, de sorte que la recette ne nécessite aucune contrainte de divisibilité.

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  # faire correspondre le nombre de tuiles à la partition

4. Définir la tâche feuille et la lancer manuellement

PATHS et CUM_ROWS (les chemins d’accès aux fichiers et la table d’offset cumulatif des lignes de l’étape 1) ainsi que TILE_ROWS sont renseignés en tant que variables globales du module par le pilote avant le lancement ; la réplication de contrôle exécute le pilote sur chaque rang, de sorte que chaque travailleur voie des valeurs identiques.

Chaque tâche construit d’abord sa vue consommateur (cupy sur GPU, numpy sur CPU/OMP) et lit le nombre réel de lignes de la tuile à partir de view.shape[0] — PhysicalStore ne dispose pas lui-même d’un attribut .shape, il est donc nécessaire de passer par la vue. Elle calcule ensuite sa plage de lignes globale à partir de sa coordonnée de lancement et de ce nombre de lignes, divise cum_rows en deux pour les fichiers qui se chevauchent, et copie chaque tranche de fichier qui se chevauche dans la tranche de destination correspondante. Enregistrer les variantes CPU, OMP et GPU afin que le même lancement s’exécute sans modification où que ce soit ; la répartition via ctx.get_variant_kind() sélectionne le consommateur correspondant à l’endroit où l’ OutputStore réside (cp.from_dlpack(dst) pour FBMEM, np.asarray(dst) pour SYSMEM). cupy n’est importé que dans la branche GPU, ce qui permet au corps de la tâche de se charger sur les machines ne disposant pas de 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]                              # index de tuile 0..num_tasks-1

    variant = ctx.get_variant_kind()
    if variant == VariantCode.GPU:
        import cupy as cp                              # paresseux : uniquement sur le GPU
        view = cp.from_dlpack(dst)
    else:
        view = np.asarray(dst)                         # vue NumPy sans copie

    tile_rows_actual = view.shape[0]                   # manque une ligne sur la dernière tuile
    row_start = t * TILE_ROWS                          # début de l'axe 0 global
    row_end = row_start + tile_rows_actual

    # Trouver l'intervalle semi-ouvert d'indices de fichiers qui chevauche [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 de la tuile [row_start, row_end) avec le fichier [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                # écriture numpy sans copie

manual_task = runtime.create_manual_task(
    load_tile.library,
    load_tile.task_id,
    (num_tasks,),                                      # nombre de domaines de lancement == nombre de tuiles
)
manual_task.add_output(partition)
manual_task.execute()

Les deux consommateurs passent par les producteurs natifs de PhysicalStore (__dlpack__ pour cupy, __array_interface__ pour np.asarray) — des vues sans copie de la tuile locale. Le coût de la bisection est de l’ordre de O(log num_shards) et la boucle interne s’exécute généralement 1 à 2 fois (les tuiles se chevauchent au maximum sur quelques fichiers).

5. Barrière et vérification

get_legate_runtime().issue_execution_fence(block=True)

Contraintes strictes

1. Tous les shards doivent partager le même dtype et les mêmes axes de fin (shape[1:]). La recette empile les shards le long de l’axe 0 ; les axes de fin de la destination proviennent de trailing_shape, que l’étape de découverte verrouille sur la valeur du premier fichier. Le nombre de lignes par fragment (shape[0]) peut varier librement — la table d’offset cumulatif les gère. L’ exemple rejette tout fragment dont le type de données ou la forme des axes finaux diffère de celui du premier, en affichant une erreur descriptive.

2. Choisissez le consommateur qui correspond à la variante. cp.from_dlpack rejette les magasins résidant en SYSMEM ; np.asarray renvoie silencieusement une vue hôte d’un magasin résidant en FBMEM sur lequel vous ne pouvez pas réellement écrire. Effectuez une répartition en fonction de ctx.get_variant_kind() afin que chaque variante utilise son propre consommateur — voir l’étape 4.

3. Les vues mmap ne sont pas toujours contiguës en C — enveloppez chaque tranche par fichier avec `np.ascontiguousarray(arr[file_lo:file_hi])` avant `.set()` ou l’écriture sur place de numpy.

4. Multi-nœuds : SHARD_DIR doit se trouver sur un système de fichiers partagé. Chaque travailleur (sur chaque rang) ouvre les fragments par chemin ; les chemins /tmp locaux au nœud ne fonctionnent que pour les démos à nœud unique.

Variantes

Parcours rapide avec des shards uniformes (une tâche par fichier)

Lorsque chaque fragment a déjà la même forme et le même type de données (shape, dtype) et que vous disposez exactement de num_shards processeurs, le mécanisme cum-rows / bisect constitue une surcharge. Définissez tile_rows = shard_shape[0] et num_tasks = num_shards; la partition comporte alors une tuile par fichier et chaque tâche lit exactement un fichier de bout en bout (pas de bisection, pas de boucle interne). Le commutateur côté pilote tient en une seule ligne :

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)

Le même corps de tâche `load_tile` fonctionne toujours dans les deux modes — la boucle interne s'exécute simplement exactement une fois par tâche. Il n’est pas nécessaire de créer un corps de tâche distinct pour le chemin rapide.

Surdécomposition pour un meilleur équilibrage de charge

La valeur par défaut `tile_rows = ceil(total_rows / num_processors) ` donne une tuile par processeur. Pour surdécomposer d’un facteur K (tuiles plus petites, tâches plus nombreuses, mise en file d’attente plus fine), divisez plutôt par `K * num_processors` :

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

num_tasks = ceil(total_rows / tile_rows) donne alors environ K * num_processors. Le même corps de tâche fonctionne toujours — la méthode de bisection aboutit simplement à un plus grand nombre de tâches par fichier.

Autres formats

Seul le lecteur par fichier à l’intérieur de `load_tile` change. Le contrat du lecteur est le suivant : étant donné un chemin de fichier et une plage de lignes semi-ouverte [file_lo, file_hi) le long de l’axe 0, renvoyer un tableau numpy de forme (file_hi - file_lo,) + trailing_shape pouvant être rendu contigu en C. Des lectures de plage/tranche peu coûteuses sont requises — les formats qui ne prennent en charge que la « lecture du fichier entier » empêchent le cas de chevauchement partiel (une tuile qui ne couvre qu’une partie d’un fichier).

(Pour les chargeurs intégrés à appel unique par format, consultez le tableau « Pourquoi cette fonctionnalité existe » en haut de ce fichier.)

L'étape de découverte (étape 1) analyse les métadonnées de chaque format : .npy / HDF5 / Parquet contiennent tous le nombre de lignes par fichier et le type de données (dtype) sur le disque. Ce n'est pas le cas des fichiers binaires bruts — il faut utiliser un fichier sidecar ou déduire ces informations à partir de la taille du fichier.

Pièges courants

cn.asarray(dst) est interdit dans une tâche feuille

À l'intérieur du corps d'une @task, toute opération cupynumeric qui interagit avec l’environnement d’exécution de niveau supérieur — cn.asarray(store), affectation par tranchon cn_dst[s] = host_np — déclenche create_index_space depuis un contexte incorrect et Legion s’interrompt :

EXCEPTION D'UTILISATION DE L'API LEGION : contexte de tâche non valide transmis à l'appel du runtime
create_index_space

Solution : utilisez la capsule DLPack avec une bibliothèque tierce (cupy / torch / numpy) à l’intérieur des tâches feuilles. cn.asarray fonctionne correctement dans le pilote, mais pas dans les tâches feuilles. Consultez le fichier examples/dlpack/leaf_task_interop.py pour découvrir la solution de contournement spécifique à torch.

Une assertion au sein d’une tâche provoque l’interruption du runtime

Legate traite les exceptions non levées dans une @task comme une violation de contrat et interrompt l’exécution, sauf si la tâche a été enregistrée avec throws_exception(). Effectuez un contrôle de cohérence sur l’hôte avant le lancement.

Le domaine de lancement doit correspondre au nombre de tuiles de partition

create_manual_task(launch_shape=...) et partition_by_tiling(...) sont indépendants — le runtime ne détecte pas de non-correspondance. Domaine de lancement plus grand → tuiles hors limites ; plus petit → tuiles non écrites. Déduisez toujours les deux à partir des mêmes (total_rows, tile_rows) via deux divisions par le plus grand entier distinctes (dimensionner le domaine de lancement directement en fonction de num_processors entraînerait un lancement excessif lorsque 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,))