Dados fragmentados em paralelo -> carregamento com cupynumeric

Por que essa habilidade existe. O cupynumeric espelha a API de matrizes do NumPy, incluindo o cupynumeric.load para um único arquivo .npy. Além disso, o carregamento de arquivos ocorre no Legate, não no cupynumeric:

Esta habilidade mostra a maneira canônica de preencher a lacuna na última linha: escreva uma tarefa Legate em Python que chame o leitor de terceiros exigido pelo formato (h5py, pyarrow, np.memmap, ...) dentro do corpo da tarefa e deixe o Legate distribuir as leituras entre GPUs/nós. Para os formatos com carregador integrado, dê preferência a ele, a menos que você precise de um corpo de tarefa personalizado (carregador baseado em mmap, decodificador específico do formato, metadados sidecar, leituras parciais/fragmentadas).

Padrão canônico: partição manual + lançamento manual da tarefa, dimensionado para a máquina, não para os arquivos. Apenas o eixo 0 é fragmentado; os eixos subsequentes são incluídos dentro de cada bloco. O número de linhas por fragmento pode variar entre os arquivos (apenas o dtype e os eixos subsequentes devem corresponder); o lançamento preenche todos os processadores disponíveis, independentemente do número de arquivos.

.npy é o exemplo prático porque o cabeçalho contém a forma e o tipo de dados no disco, mas a estrutura se aplica a qualquer formato com leituras de intervalo/fatia de baixo custo (binário bruto, HDF5, Parquet/Arrow — veja “Outros formatos” abaixo). Implementação de referência: assets/examples/parallel_npy_load.py.

Suposição sobre o layout dos dados

Esta habilidade trata exclusivamente do carregamento — ela pressupõe que os dados já estejam dispostos em um sistema de arquivos compartilhado de alguma forma previsível e indexável. A geração desses arquivos está fora do escopo (o exemplo inclui um subcomando de gravação por conveniência, mas os usuários reais devem usar os seus próprios).

O exemplo prático pressupõe um layout específico:

• Um diretório contendo arquivos nomeados shard_0000.npy, shard_0001.npy, ... em uma sequência contígua de números inteiros (largura de 4 preenchida com zeros).
• Todos os fragmentos compartilham o mesmo dtype e os mesmos eixos finais (shape[1:]); o eixo 0 (linhas por fragmento) pode variar entre os arquivos — a receita constrói uma tabela cumulativa de deslocamento de linha e lê a fatia sobreposta de cada arquivo a partir da tarefa folha.
• O diretório é visível para todos os ranks (sistema de arquivos compartilhado para execuções com múltiplos nós).

A função discover_layout() do exemplo exibe o que encontrou e apresenta uma falha grave com um erro descritivo quando o layout está incorreto (diretório ausente, ausência de shards, incompatibilidade de dtype/ eixos finais ou uma lacuna na sequência contígua shard_NNNN.npy ).

Se seus dados estiverem em um layout diferente — binário bruto de passo fixo, um arquivo HDF5 com um conjunto de dados por fragmento, uma árvore de diretórios, ... — apenas o padrão glob, o leitor por arquivo (etapa 4 abaixo) e a descoberta de metadados (etapa 1 abaixo) mudam. O mecanismo de particionamento e inicialização é independente do layout.

Quando usar

Consulte a tabela de formatos acima para a decisão de roteamento (carregador embutido vs. esta skill). Além disso, há dois indícios adicionais de que esta skill é a escolha certa:

• Substituir o código sequencial `np.concatenate([read(f) for f in files]) ` por leituras paralelas por GPU.
• Demonstrar como uma tarefa Legate Python definida pelo usuário grava em um matriz de saída cupynumeric por meio de uma execução manual.

Exemplos

Os caminhos abaixo são escritos em relação ao diretório desta habilidade (o script vem em assets/examples/parallel_npy_load.py). Ajuste o prefixo para corresponder ao local onde sua habilidade está instalada (por exemplo, skills/cupynumeric-parallel-data-load/assets/... se a habilidade estiver no diretório de nível superior skills/ ).

# Um único nó, 4 GPUs.
legate --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

# Multinó, 2 nós x 4 GPUs (slurm), sistema de arquivos compartilhado em --shard-dir.
# Gere os fragmentos uma vez no rank 0 e, em seguida, execute novamente `read` em qualquer 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

Sem sinalizadores de layout — o driver de leitura percorre cada cabeçalho .npy para recuperar a contagem de linhas por arquivo, a forma final e o tipo de dados (dtype), e então calcula tile_rows a partir da contagem de processadores disponíveis.

--min-gpu-chunk 1 só é necessário quando a contagem de elementos por bloco está abaixo do tamanho mínimo padrão de bloco do Legate para execuções em GPU (por exemplo, os padrões do exemplo prático — total de linhas divididas entre 4 GPUs a cerca de 1 milhão por bloco — ficam abaixo do limite e, caso contrário, seriam agrupados em uma única GPU). Para conjuntos de dados de tamanho de produção (dezenas de milhões de elementos por bloco ou mais), você pode omitir a opção e deixar o Legate usar seu padrão. Aumentá-la para um valor moderado (por exemplo, --min-gpu-chunk 1024) é adequado quando cada bloco é grande o suficiente para que a sobrecarga por tarefa seja mais importante do que atribuir um bloco a cada GPU.

Instruções

Cinco etapas a partir de um exemplo prático em .npy; apenas a etapa 1 (análise do cabeçalho do formato) e a etapa 4 (o leitor por arquivo dentro do corpo da tarefa) são específicas do formato. Os outros três (alocação de destino, partição, barreira) são reutilizados sem alterações em todos os formatos — consulte “Outros formatos” abaixo para os pontos de troca.

1. Leia os metadados de cada fragmento

Faça uma varredura no diretório e examine cada cabeçalho .npy (mmap_mode="r" lê apenas o cabeçalho). O cabeçalho contém a forma e o tipo de dados por fragmento, de modo que o driver pode recuperar o total de linhas, a forma final e uma tabela cumulativa de deslocamento de linhas sem precisar carregar os dados:

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

per_file_rows = []                       # linhas ao longo do eixo 0 por arquivo
trailing_shape = None                    # shape[1:], deve corresponder entre os arquivos
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}: incompatibilidade entre a forma final e o tipo de dados "
            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)  # comprimento N+1
total_rows = int(cum_rows[-1])

O trecho de código acima garante a correspondência entre dtype e trailing_shape (ou seja, shape[1:]) entre os arquivos. A contagem de linhas por fragmento pode variar — a tabela cum_rows lida com isso. O código de produção também deve verificar se os nomes formam uma sequência contígua de shard_0000.npy ... shard_NNNN.npy (omitida do trecho por uma questão de concisão; consulte discover_layout() no exemplo prático). A descoberta depende apenas do que o próprio formato no disco expõe (o cabeçalho .npy aqui, .shape / .dtype para HDF5, etc.); qualquer arquivo auxiliar (manifesto, hashes de conteúdo) é uma etapa de verificação separada adicional.

2. Crie o armazenamento de saída cupynumeric a partir dos metadados

A matriz total abrange total_rows ao longo do eixo 0; os eixos restantes vêm de trailing_shape inalterados. Use cn.empty — a tarefa sobrescreve todas as células, a inicialização com zeros seria um desperdício.

import cupynumeric as cn

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

3. Divida o armazenamento em blocos de acordo com o número de processadores

A forma de inicialização é dimensionada de acordo com os processadores disponíveis, não com a contagem de arquivos. Escolha tile_rows = ceil(total_rows / num_processors) e particione o eixo 0 por esse tamanho de bloco. Os eixos finais não são particionados (o bloco abrange toda a extensão nesse eixo). É permitido que o último bloco seja curto — é exatamente isso que o `partition_by_tiling` suporta — portanto, a receita não precisa de nenhuma restrição de divisibilidade.

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  # fazer corresponder a contagem de blocos à partição

4. Defina a tarefa folha e inicie-a manualmente

PATHS e CUM_ROWS (os caminhos dos arquivos e a tabela de deslocamento cumulativo de linhas da etapa 1), além de TILE_ROWS, são preenchidos como variáveis globais do módulo pelo driver antes do lançamento; o controle de replicação executa o driver em cada rank, de modo que cada trabalhador veja valores idênticos.

Cada tarefa constrói primeiro sua visão de consumidor (cupy na GPU, numpy na CPU/OMP) e lê a contagem real de linhas do bloco a partir de view.shape[0] — o próprio PhysicalStore não possui o atributo .shape, portanto, é necessário passar pela visão. Em seguida, ela calcula seu intervalo global de linhas a partir de sua coordenada de inicialização e dessa contagem de linhas, divide cum_rows ao meio para os arquivos sobrepostos e copia cada fatia de arquivo sobreposta para a fatia de destino correspondente. Registre as variantes de CPU, OMP e GPU para que a mesma inicialização seja executada sem alterações em qualquer lugar; o despacho por meio de `ctx.get_variant_kind()` seleciona o consumidor que corresponde ao local onde o OutputStore está residente (cp.from_dlpack(dst) para FBMEM, np.asarray(dst) para SYSMEM). O cupy é importado apenas dentro do ramo da GPU, portanto, o corpo da tarefa é carregado em máquinas sem 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 do bloco 0..num_tasks-1

    variant = ctx.get_variant_kind()
    if variant == VariantCode.GPU:
        import cupy as cp                              # preguiçoso: somente na GPU
        view = cp.from_dlpack(dst)
    else:
        view = np.asarray(dst)                         # visualização numpy sem cópia

    tile_rows_actual = view.shape[0]                   # falta uma linha no último bloco
    row_start = t * TILE_ROWS                          # início global no eixo 0
    row_end = row_start + tile_rows_actual

    # Encontrar o intervalo semiaberto de índices de arquivo que se sobrepõem a [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):
        # Interseção do bloco [row_start, row_end) com o arquivo [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                # gravação numpy sem cópia

manual_task = runtime.create_manual_task(
    load_tile.library,
    load_tile.task_id,
    (num_tasks,),                                      # número de tarefas lançadas == número de blocos
)
manual_task.add_output(partition)
manual_task.execute()

Ambos os consumidores passam pelos produtores nativos do PhysicalStore (__dlpack__ para cupy, __array_interface__ para np.asarray) — visualizações sem cópia do bloco local. O custo da bissecção é O(log num_shards) e o loop interno normalmente itera de 1 a 2 vezes (os blocos se sobrepõem, no máximo, em alguns arquivos).

5. Fence e verificação

get_legate_runtime().issue_execution_fence(block=True)

Restrições rígidas

1. Todos os fragmentos devem compartilhar o dtype e os eixos finais (shape[1:]). A receita empilha os fragmentos ao longo do eixo 0; os eixos finais do destino provêm de trailing_shape, que a etapa de descoberta bloqueia no valor do primeiro arquivo. O número de linhas por fragmento (shape[0]) pode variar livremente — a tabela de deslocamento cumulativo lida com isso. O exemplo rejeita qualquer fragmento cujo dtype ou forma final difira do primeiro, exibindo um erro descritivo.

2. Escolha o consumidor que corresponda à variante. cp.from_dlpack rejeita armazenamentos residentes em SYSMEM; np.asarray retorna silenciosamente uma visão do host de um armazenamento residente em FBMEM no qual você não pode, na verdade, gravar . Faça o despacho com base em ctx.get_variant_kind() para que cada variante utilize seu próprio consumidor — consulte a etapa 4.

3. As visualizações mmap nem sempre são contíguas em C — envolva cada fatia por arquivo com np.ascontiguousarray(arr[file_lo:file_hi]) antes de .set() ou da gravação in-place do numpy.

4. Multinó: SHARD_DIR deve estar em um sistema de arquivos compartilhado. Cada trabalhador (em cada rank) abre os fragmentos pelo caminho; caminhos /tmp locais ao nó funcionam apenas para demonstrações de nó único.

Variantes

Caminho rápido com shards uniformes (uma tarefa por arquivo)

Quando cada fragmento já possui a mesma forma e tipo de dados (shape, dtype) e você por acaso tiver num_shards processadores disponíveis, o mecanismo de cum-rows / bisect representa uma sobrecarga. Defina tile_rows = shard_shape[0] e num_tasks = num_shards; a partição passa a ter um bloco por arquivo e cada tarefa lê exatamente um arquivo do início ao fim (sem bisect, sem loop interno). A alteração no driver é feita em uma única linha:

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)

O mesmo corpo da tarefa `load_tile ` continua funcionando em ambos os modos — o loop interno simplesmente itera exatamente uma vez por tarefa. Não há necessidade de um corpo de tarefa separado para o caminho rápido.

Superdecompor para um melhor balanceamento de carga

O padrão tile_rows = ceil(total_rows / num_processors) resulta em um bloco por processador. Para decomposição excessiva por um fator K (blocos menores, mais tarefas pontuais, enfileiramento mais refinado), divida por K * num_processors em vez disso:

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

num_tasks = ceil(total_rows / tile_rows) então se expande para aproximadamente K * num_processors. O mesmo corpo de tarefa ainda funciona — a bissecção simplesmente resulta em mais tarefas por arquivo.

Outros formatos

Apenas o leitor por arquivo dentro de `load_tile` muda. O contrato do leitor: dado um caminho de arquivo e um intervalo de linhas semiaberto [file_lo, file_hi) ao longo do eixo 0, retorne uma matriz numpy com a forma (file_hi - file_lo,) + trailing_shape que possa ser tornado contíguo em C. São necessárias leituras baratas de intervalos/fatias — formatos que suportam apenas “ler o arquivo inteiro” inviabilizam o caso de sobreposição parcial (um bloco que cobre apenas parte de um arquivo).

(Para carregadores embutidos de chamada única por formato, consulte a tabela “Por que esta habilidade existe” no início deste arquivo.)

A etapa de descoberta (etapa 1) analisa os metadados de cada formato: .npy / HDF5 / Parquet contêm, no disco, a contagem de linhas por arquivo + o dtype. O binário bruto não contém — é necessário usar um arquivo sidecar ou derivar a partir do tamanho do arquivo.

Armadilhas comuns

cn.asarray(dst) é inválido em uma tarefa folha

Dentro do corpo de uma @task, qualquer operação cupynumeric que acesse o tempo de execução de nível superior — cn.asarray(store), atribuição de fatia cn_dst[s] = host_np — aciona create_index_space a partir do contexto errado e o Legion é abortado:

EXCEÇÃO DE USO DA API DO LEGION: Contexto de tarefa inválido passado para a chamada de tempo de execução
create_index_space

Solução: utilize a cápsula DLPack com uma biblioteca de terceiros (cupy / torch / numpy) dentro de tarefas folha. cn.asarray funciona normalmente no driver, mas não em tarefas folha. Consulte examples/dlpack/leaf_task_interop.py para ver a solução alternativa com o torch.

Uma assert dentro da tarefa interrompe o tempo de execução

O Legate trata exceções não levantadas em uma @task como uma violação de contrato e interrompe, a menos que a tarefa tenha sido registrada com throws_exception(). Faça uma verificação de integridade no host antes de iniciar.

O domínio de inicialização deve corresponder ao número de blocos de partição

create_manual_task(launch_shape=...) e partition_by_tiling(...) são independentes — o tempo de execução não detecta uma incompatibilidade. Domínio de lançamento maior → blocos fora do intervalo; menor → blocos não gravados. Sempre derive ambos a partir do mesmo (total_rows, tile_rows) por meio de duas divisões por teto separadas (dimensionar o domínio de lançamento diretamente para num_processors resultaria em lançamento excessivo quando 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,))