Datos fragmentados en paralelo -> carga con cupynumeric

Por qué existe esta habilidad. cupynumeric reproduce la API de matrices de NumPy, incluido cupynumeric.load para un único archivo .npy. Más allá de eso, la carga de archivos se realiza en Legate, no en cupynumeric:

Esta habilidad muestra la forma canónica de cubrir el hueco de la última fila: escribe una tarea de Legate en Python que llame al lector de terceros que requiera el formato (h5py, pyarrow, np.memmap, …) dentro del cuerpo de la tarea, y deja que Legate distribuya las lecturas entre las GPU o los nodos. Para los formatos con un cargador integrado, utilízalo siempre que sea posible, a menos que necesites un cuerpo de tarea personalizado (cargador basado en mmap, decodificador específico del formato, metadatos sidecar, lecturas parciales o fragmentadas).

Patrón canónico: partición manual + lanzamiento manual de la tarea, dimensionada en función de la máquina, no de los archivos. Solo se fragmenta el eje 0; los ejes restantes se incluyen dentro de cada mosaico. El recuento de filas por fragmento puede variar entre archivos (solo deben coincidir el tipo de datos y los ejes restantes); el lanzamiento ocupa todos los procesadores disponibles, independientemente del número de archivos.

.npy es el ejemplo práctico porque el encabezado contiene la forma y el tipo de datos en el disco, pero el esqueleto se aplica a cualquier formato con lecturas de rango/segmento económicas (binario sin procesar, HDF5, Parquet/Arrow; véase «Otros formatos» más abajo). Implementación de referencia: assets/examples/parallel_npy_load.py.

Supuesto sobre la estructura de los datos

Esta función se centra exclusivamente en la carga: da por hecho que los datos ya están dispuestos en un sistema de archivos compartido de forma predecible e indexable. La creación de esos archivos queda fuera del alcance de esta función (el ejemplo incluye un subcomando de escriturapor comodidad, pero los usuarios reales deben aportar el suyo propio).

El ejemplo práctico asume una estructura específica:

• Un directorio que contiene archivos denominados shard_0000.npy, shard_0001.npy, ... en una secuencia contigua de números enteros (anchura de 4, rellenada con ceros).
• Todos los fragmentos comparten el mismo tipo de datos (dtype) y los mismos ejes finales (shape[1:]); el eje 0 (filas por fragmento) puede variar entre los archivos — la receta crea una tabla acumulativa de desplazamiento de filas y lee el segmento superpuesto de cada archivo desde el interior de la tarea hoja.
• El directorio es visible para todos los rangos (sistema de archivos compartido para ejecuciones con varios nodos).

La función discover_layout() del ejemplo muestra lo que ha encontrado y devuelve un error grave con un mensaje descriptivo cuando la estructura es incorrecta (falta el directorio, no hay fragmentos, el tipo de datos o los ejes finales no coinciden, o hay un hueco en la secuencia contigua shard_NNNN.npy ).

Si tus datos tienen una estructura diferente —binario sin procesar de paso fijo, un archivo HDF5 con un conjunto de datos por fragmento, un árbol de directorios, etc.—, solo cambian el patrón glob, el lector por archivo (paso 4 más abajo) y la detección de metadatos (paso 1 más abajo). El mecanismo de partición y ejecución es independiente del formato de los datos.

Cuándo utilizarla

Consulta la tabla de formatos anterior para decidir el enrutamiento (cargador integrado frente a esta habilidad). Además, hay dos indicios adicionales de que esta habilidad es la opción adecuada:

• Sustituir el código secuencial `np.concatenate([read(f) for f in files])` por lecturas paralelas por GPU.
• Demostrar cómo una tarea Legate Python definida por el usuario escribe en una matriz de salida cupynumeric mediante un lanzamiento manual.

Ejemplos

Las rutas que se indican a continuación están escritas de forma relativa al directorio de esta habilidad (el script se incluye en assets/examples/parallel_npy_load.py). Ajusta el prefijo para que coincida con la ubicación en la que esté instalada tu habilidad (por ejemplo, skills/cupynumeric-parallel-data-load/assets/... si la habilidad se encuentra en un directorio de nivel superior llamado skills/ ).

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

# Multinodo, 2 nodos x 4 GPU (slurm), sistema de archivos compartido en --shard-dir.
# Genera los fragmentos una vez en el rango 0 y, a continuación, vuelve a ejecutar `read` a cualquier escala.
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

Sin indicadores de disposición: el controlador de lectura recorre cada encabezado .npy para recuperar el recuento de filas por archivo, la forma final y el tipo de datos, y a continuación calcula tile_rows a partir del número de procesadores disponibles.

--min-gpu-chunk 1 solo es necesario cuando el recuento de elementos por mosaico es inferior al tamaño mínimo predeterminado de Legate para los lanzamientos en GPU (por ejemplo, los valores predeterminados del ejemplo práctico: el total de filas repartidas entre 4 GPU a ~1 millón por mosaico— caen por debajo del umbral y, de no ser así, se agruparían en una sola GPU). Para conjuntos de datos de tamaño de producción (decenas de millones de elementos por mosaico o más), puedes omitir la opción y dejar que Legate utilice su valor por defecto. Aumentarlo a un valor moderado (p. ej., --min-gpu-chunk 1024) es adecuado cuando cada bloque es lo suficientemente grande como para que la sobrecarga por tarea sea más importante que asignar un bloque a cada GPU.

Instrucciones

Cinco pasos a partir de un ejemplo práctico en formato .npy; solo el paso 1 (análisis del encabezado del formato) y el paso 4 (el lector por archivo dentro del cuerpo de la tarea) son específicos del formato. Los otros tres (asignación de destino, partición y barrera) se reutilizan sin cambios en todos los formatos; consulta «Otros formatos» más abajo para conocer los puntos de intercambio.

1. Leer los metadatos de cada fragmento

Explora el directorio y echa un vistazo a cada encabezado .npy (mmap_mode="r" solo lee el encabezado). El encabezado contiene la forma y el tipo de datos por fragmento, por lo que el controlador puede recuperar el número total de filas, la forma final y una tabla acumulativa de desplazamientos de filas sin necesidad de cargar los datos:

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

per_file_rows = []                       # filas a lo largo del eje 0 por archivo
trailing_shape = None                    # shape[1:], debe coincidir en todos los archivos
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}: discrepancia entre la forma final y el tipo de datos "
            f"({hdr.shape[1:]}/{hdr.dtype} frente a {trailing_shape}/{dtype})"
        )
    per_file_rows.append(int(hdr.shape[0]))

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

El fragmento de código anterior garantiza que el tipo de datos (dtype) y la forma final (trailing_shape) (es decir, shape[1:]) coincidan en todos los archivos. El recuento de filas por fragmento puede variar; la tabla cum_rows se encarga de ello. El código de producción también debe verificar que los nombres formen una secuencia contigua shard_0000.npy ... shard_NNNN.npy (omitida en el fragmento por brevedad; véase discover_layout() en el ejemplo práctico). La detección se basa únicamente en lo que el propio formato en disco expone (la cabecera .npy en este caso, .shape / .dtype para HDF5, etc.); cualquier archivo complementario (manifiesto, hash de contenido) supone un paso de verificación adicional.

2. Crear el almacén de salida cupynumeric a partir de los metadatos

La matriz total abarca total_rows a lo largo del eje 0; los ejes restantes provienen de trailing_shape sin modificaciones. Utiliza cn.empty: la tarea sobrescribe todas las celdas, por lo que una inicialización a cero sería innecesaria.

import cupynumeric as cn

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

3. Dividir el almacenamiento en mosaicos según el número de procesadores

La forma de lanzamiento se adapta al número de procesadores disponibles, no al número de archivos. Elige tile_rows = ceil(total_rows / num_processors) y divide el eje 0 según ese tamaño de mosaico. Los ejes restantes no se dividen (el mosaico abarca toda la extensión en ese caso). Se permite que el último mosaico sea más corto —eso es precisamente lo que admite partition_by_tiling—, por lo que la receta no necesita ninguna restricción de divisibilidad.

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  # hacer coincidir el número de mosaicos con la partición

4. Definir la tarea hoja y ejecutarla manualmente

PATHS y CUM_ROWS (las rutas de los archivos y la tabla de desplazamiento acumulativo de filas del paso 1), junto con TILE_ROWS, se rellenan como variables globales del módulo por el controlador antes de la ejecución; el control de replicación ejecuta el controlador en cada rango, de modo que todos los trabajadores ven valores idénticos.

Cada tarea crea primero su vista de consumidor (cupy en la GPU, numpy en la CPU/OMP) y lee el recuento real de filas del mosaico a partir de view.shape[0] — PhysicalStore en sí mismo no tiene el atributo .shape, por lo que es necesario acceder a través de la vista. A continuación, calcula su rango global de filas a partir de su coordenada de inicio y ese recuento de filas, divide cum_rows por la mitad para los archivos superpuestos y copia cada segmento de archivo superpuesto en el segmento de destino correspondiente. Se registran las variantes de CPU, OMP y GPU para que el mismo lanzamiento se ejecute sin cambios en cualquier lugar; el envío mediante ctx.get_variant_kind() selecciona el consumidor que coincida con el lugar donde el resideel OutputStore (cp.from_dlpack(dst) para FBMEM, np.asarray(dst) para SYSMEM). cupy se importa únicamente en la rama de la GPU, por lo que el cuerpo de la tarea se carga en máquinas sin 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]                              # índice de mosaico 0..num_tasks-1

    variant = ctx.get_variant_kind()
    if variant == VariantCode.GPU:
        import cupy as cp                              # diferido: solo en GPU
        view = cp.from_dlpack(dst)
    else:
        view = np.asarray(dst)                         # vista numpy sin copia

    tile_rows_actual = view.shape[0]                   # falta una fila en el último mosaico
    row_start = t * TILE_ROWS                          # inicio global del eje 0
    row_end = row_start + tile_rows_actual

    # Hallar el intervalo semiabierto de índices de archivo que se solapan con [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):
        # Intersección del mosaico [row_start, row_end) con el archivo [cum[f], cum[f+1]).
        lo = max(row_start, int(CUM_ROWS[f]))
        hi = min(fin_fila, int(CUM_ROWS[f + 1]))
        límite_inferior_ficha = lo - int(CUM_ROWS[f])
        límite_superior_ficha = 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                # escritura numpy sin copia

manual_task = runtime.create_manual_task(
    load_tile.library,
    load_tile.task_id,
    (num_tasks,),                                      # número de dominios de ejecución == número de mosaicos
)
manual_task.add_output(partition)
manual_task.execute()

Ambos consumidores pasan por los productores nativos de PhysicalStore (__dlpack__ para cupy, __array_interface__ para np.asarray): vistas sin copia del mosaico local. El coste de la bisección es O(log num_shards) y el bucle interno suele iterar entre 1 y 2 veces (los mosaicos se solapan como máximo en un par de archivos).

5. Barrera y verificación

get_legate_runtime().issue_execution_fence(block=True)

Restricciones estrictas

1. Todos los fragmentos deben compartir tipo de datos (dtype) y ejes finales (shape[1:]). La receta apila los fragmentos a lo largo del eje 0; los ejes finales del destino proceden de trailing_shape, que la fase de descubrimiento fija en el valor del primer archivo. El número de filas por fragmento (shape[0]) puede variar libremente: la tabla de desplazamiento acumulativo se encarga de gestionarlas. El ejemplo rechaza cualquier fragmento cuyo tipo de datos (dtype) o forma final difiera del primero, mostrando un error descriptivo.

2. Elige el consumidor que coincida con la variante. cp.from_dlpack rechaza los almacenes residentes en SYSMEM; np.asarray devuelve silenciosamente una vista del host de un almacén residente en FBMEM en el que, en realidad, no se puede escribir . Realiza la distribución según ctx.get_variant_kind() para que cada variante utilice su propio consumidor — véase el paso 4.

3. Las vistas mmap no siempre son contiguas en C: envuelve cada segmento por archivo con np.ascontiguousarray(arr[file_lo:file_hi]) antes de .set() o de la escritura in situ de numpy.

4. Multinodo: SHARD_DIR debe estar en un sistema de archivos compartido. Cada trabajador (en cada rango) abre los fragmentos por ruta; las rutas /tmp locales del nodo solo funcionan para demostraciones de un solo nodo.

Variantes

Ruta rápida con fragmentos uniformes (una tarea por archivo)

Cuando todos los fragmentos tienen ya la misma forma y tipo de datos (shape, dtype) y se dispone de num_shards procesadores, el mecanismo de cum-rows / bisect supone una sobrecarga. Establece tile_rows = shard_shape[0] y num_tasks = num_shards; la partición tendrá entonces un mosaico por archivo y cada tarea leerá exactamente un archivo de principio a fin (sin bisección, sin bucle interno). El cambio en el controlador se realiza en una sola línea:

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)

El mismo cuerpo de la tarea `load_tile ` sigue funcionando en ambos modos; el bucle interno simplemente itera exactamente una vez por tarea. No es necesario un cuerpo de tarea independiente para la ruta rápida.

Descomponer en exceso para un mejor equilibrio de carga

El valor predeterminado tile_rows = ceil(total_rows / num_processors) asigna un mosaico por procesador. Para sobredescomponer en un factor K (mosaicos más pequeños, más tareas puntuales, colas más detalladas), divide por K * num_processors en su lugar:

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

num_tasks = ceil(total_rows / tile_rows) se amplía entonces hasta aproximadamente K * num_processors. El mismo cuerpo de la tarea sigue funcionando: la bisección simplemente da como resultado un mayor número de tareas por archivo.

Otros formatos

Solo cambia el lector por archivo dentro de `load_tile`. El contrato del lector: dada una ruta de archivo y un rango de filas semiabierto [file_lo, file_hi) a lo largo del eje 0, devuelve una matriz numpy de forma (file_hi - file_lo,) + trailing_shape que pueda hacerse contiguo en C. Se requieren lecturas de rango/segmento económicas: los formatos que solo admiten «leer todo el archivo» impiden el caso de solapamiento parcial (un mosaico que solo cubre parte de un archivo).

(Para conocer los cargadores integrados de una sola llamada por formato, consulta la tabla «Por qué existe esta función» que se encuentra al principio de este archivo).

El paso de descubrimiento (paso 1) analiza los metadatos de cada formato: .npy / HDF5 / Parquet incluyen en el disco el recuento de filas por archivo y el tipo de datos (dtype). El formato binario sin procesar no los incluye; se deducen del tamaño del archivo o mediante un archivo auxiliar.

Errores habituales

cn.asarray(dst) no es válido en una tarea hoja

Dentro del cuerpo de una @task, cualquier operación cupynumeric que interfiera con el tiempo de ejecución de nivel superior — cn.asarray(store), asignación de corte cn_dst[s] = host_np — desencadena create_index_space desde un contexto incorrecto y Legion se interrumpe:

EXCEPCIÓN DE USO DE LA API DE LEGION: Se ha pasado un contexto de tarea no válido a la llamada de tiempo de ejecución
create_index_space

Solución: utiliza la cápsula DLPack con una biblioteca de terceros (cupy / torch / numpy) dentro de las tareas hoja. cn.asarray funciona correctamente en el controlador, pero no en las tareas hoja. Consulta examples/dlpack/leaf_task_interop.py para ver la solución alternativa basada en torch.

Una afirmación dentro de la tarea interrumpe el tiempo de ejecución

Legate trata las excepciones no lanzadas en una @task como una violación del contrato y la interrumpe, a menos que la tarea se haya registrado con throws_exception(). Realiza una comprobación de validación en el host antes de iniciarla.

El dominio de lanzamiento debe coincidir con el recuento de mosaicos de partición

create_manual_task(launch_shape=...) y partition_by_tiling(...) son independientes: el tiempo de ejecución no detecta una discrepancia. Un dominio de lanzamiento más grande → mosaicos fuera de rango; más pequeño → mosaicos sin escribir. Deriva siempre ambos a partir del mismo (total_rows, tile_rows) mediante dos divisiones por el techo separadas (ajustar el tamaño del dominio de lanzamiento directamente a num_processors provocaría un lanzamiento excesivo cuando 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,))