Параллельная загрузка данных с фрагментацией -> cupynumeric

Зачем нужен этот навык. cupynumeric дублирует API массивов NumPy, включая cupynumeric.load для одного файла .npy. Помимо этого, загрузка файлов осуществляется в Legate, а не в cupynumeric:

Этот навык демонстрирует канонический способ восполнить пробел в последней строке: напишите задачу Legate на Python, которая вызывает сторонний считыватель, необходимый для данного формата (h5py, pyarrow, np.memmap и т. д.), внутри тела задачи, и позвольте Legate распределить чтение между GPU / узлами. Для форматов со встроенным загрузчиком отдавайте ему предпочтение, если вам не требуется пользовательский код внутри тела задачи (загрузчик на основе mmap, формат-специфический декодер, метаданные sidecar, частичное / фрагментированное чтение).

Стандартный шаблон: ручное разделение + ручной запуск задачи, размер которой подбирается исходя из возможностей машины, а не файлов. Разбивается только ось 0; последующие оси перемещаются внутри каждого фрагмента. Количество строк в каждом фрагменте может различаться между файлами (соответствовать должны только dtype и последующие оси); при запуске заполняются все доступные процессоры, независимо от количества файлов.

Файл .npy служит примером, поскольку его заголовок содержит информацию о форме и типе данных на диске, но эта структура применима к любому формату с недорогим чтением диапазонов/фрагментов (необработанные двоичные данные, HDF5, Parquet/Arrow — см. раздел «Другие форматы» ниже). Эталонная реализация: assets/examples/parallel_npy_load.py.

Предположение о расположении данных

Этот навык касается исключительно загрузки — он предполагает, что данные уже расположены в общей файловой системе каким-либо предсказуемым и индексируемым образом. Создание таких файлов выходит за рамки данного руководства (в примере для удобства прилагается подкоманда записи, но реальные пользователи используют свои собственные).

В примере рассматривается одна конкретная организация данных:

• каталог, содержащий файлы с именами shard_0000.npy, shard_0001.npy, ... в виде непрерывной целочисленной последовательности (ширина 4 с нулевым заполнением).
• Все фрагменты имеют одинаковый тип данных (dtype) и одинаковые конечные оси (shape[1:]); ось 0 (количество строк на фрагмент) может различаться между файлами — рецепт строит кумулятивную таблицу смещений строк и считывает перекрывающийся фрагмент каждого файла изнутри листьевой задачи.
• Каталог доступен для каждого ранга (общая файловая система для многоузловых запусков).

Функция discover_layout() в примере выводит найденные данные и завершает работу с ошибкой с подробным описанием, если расположение данных неверно (отсутствует каталог, нет фрагментов, несовпадение dtype / конечных осей или пробел в непрерывной последовательности shard_NNNN.npy ).

Если ваши данные имеют иную структуру — необработанные двоичные данные с фиксированным шагом, файл HDF5 с одним набором данных на каждый фрагмент, дерево каталогов и т. д. — изменятся только шаблон glob, считыватель для каждого файла (шаг 4 ниже) и обнаружение метаданных (шаг 1 ниже). Механизм разбиения на части и запуска не зависит от структуры данных.

Когда использовать

См. таблицу форматов выше для принятия решения о маршрутизации (встроенный загрузчик или данный навык). Помимо этого, есть ещё два дополнительных признака того, что данный навык подходит:

• Замена последовательного np.concatenate([read(f) for f in files]) на параллельное чтение с отдельным доступом к каждому GPU.
• Демонстрация того, как пользовательская задача Legate на Python записывает данные в массив вывода cupynumeric посредством ручного запуска.

Примеры

Приведённые ниже пути указаны относительно каталога данного навыка (скрипт находится в assets/examples/parallel_npy_load.py). Измените префикс в соответствии с местом установки вашего навыка (например, skills/cupynumeric-parallel-data-load/assets/... если навык находится в каталоге верхнего уровня skills/ ).

# Один узел, 4 графических процессора.
legate --gpus 4 --fbmem 4000 --min-gpu-chunk 1 \
    assets/examples/parallel_npy_load.py \
    read --shard-dir /shared/scratch/demo

# Многоузловая конфигурация, 2 узла по 4 графических процессора (slurm), общая файловая система в --shard-dir.
# Сгенерируйте фрагменты один раз на ранке 0, затем запустите `read` в любом масштабе.
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 для восстановления количества строк в каждом файле, конечной формы и типа данных, а затем вычисляет tile_rows на основе доступного количества процессоров.

Флаг --min-gpu-chunk 1 необходим только в том случае, если количество элементов в каждой плитке меньше минимального размера блока по умолчанию Legate для запусков на GPU (например, по умолчанию в примере — общее количество строк распределяется между 4 GPU по ~1 млн на каждый тайл — опускаются ниже порогового значения и в противном случае были бы объединены на одном GPU). Для наборов данных производственного масштаба (десятки миллионов элементов на тайл или больше) можно отказаться от этого флага и позволить Legate использовать значение по умолчанию. Увеличение этого значения до умеренного (например, --min-gpu-chunk 1024) целесообразно, когда каждый фрагмент достаточно велик, чтобы накладные расходы на каждую задачу имели большее значение, чем распределение фрагментов по каждому GPU.

Инструкции

Пять шагов из рабочего примера в формате .npy; только шаг 1 (разбор заголовка формата) и шаг 4 (считыватель по файлам внутри тела задачи) зависят от конкретного формата. Остальные три (выделение места назначения, разбиение на фрагменты, барьер) повторно используются без изменений во всех форматах — см. раздел «Другие форматы» ниже для точек переключения.

1. Считывание метаданных из каждого фрагмента

Просканируйте каталог и просмотрите каждый заголовок .npy (mmap_mode="r" считывает только заголовок). Заголовок содержит форму и тип данных для каждого фрагмента, поэтому драйвер может восстановить общее количество строк, форму конечных данных и таблицу совокупных смещений строк, даже не загружая данные:

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

per_file_rows = []                       # количество строк по оси 0 в каждом файле
trailing_shape = None                    # shape[1:], должно совпадать во всех файлах
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}: несовпадение формы / типа данных"
            f"({hdr.shape[1:]}/{hdr.dtype} против {trailing_shape}/{dtype})"
        )
    per_file_rows.append(int(hdr.shape[0]))

cum_rows = np.cumsum([0] + per_file_rows, dtype=np.int64)  # длина N+1
total_rows = int(cum_rows[-1])

Приведённый выше фрагмент кода обеспечивает соответствие dtype и trailing_shape (т. е. shape[1:]) между файлами. Количество строк в каждом шарде может различаться — с этим справляется таблица cum-rows. В рабочем коде также следует проверить, что имена образуют непрерывную последовательность shard_0000.npy … shard_NNNN.npy (опущено в фрагменте для краткости; см. discover_layout() в примере с решением). Обнаружение опирается исключительно на то, что предоставляет сам формат на диске (в данном случае заголовок .npy, .shape / .dtype для HDF5 и т. д.); любые дополнительные элементы (манифест, хэши содержимого) являются отдельным этапом проверки сверху.

2. Создание выходного хранилища cupynumeric на основе метаданных

Общий массив охватывает total_rows по оси 0; последующие оси берутся из 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  # сопоставить количество тайлов с количеством задач

4. Определите конечную задачу и запустите её вручную

PATHS и CUM_ROWS (пути к файлам и таблица кумулятивных смещений строк из шага 1) плюс TILE_ROWS заполняются драйвером в качестве глобальных переменных модуля перед запуском; управление репликацией запускает драйвер на каждом ранке, поэтому каждый рабочий узел видит одинаковые значения.

Каждая задача сначала создаёт своё представление-потребитель (cupy на GPU, numpy на CPU/OMP) и считывает фактическое количество строк тайла из view.shape[0] — у самого PhysicalStore нет атрибута .shape, поэтому необходимо пройти через представление. Затем она вычисляет свой глобальный диапазон строк на основе координаты запуска и этого количества строк, делит cum_rows пополам для перекрывающихся файлов и копирует каждый перекрывающийся фрагмент файла в соответствующий фрагмент назначения. Зарегистрируйте варианты для CPU, OMP и GPU, чтобы один и тот же запуск выполнялся без изменений в любом месте; диспетчеризация по ctx.get_variant_kind() выбирает потребителя, соответствующего тому, где находитсяOutputStore (cp.from_dlpack(dst) для FBMEM, np.asarray(dst) для SYSMEM). 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]                              # индекс тайла 0..num_tasks-1

    variant = ctx.get_variant_kind()
    if variant == VariantCode.GPU:
        import cupy as cp                              # отложенная загрузка: только на GPU
        view = cp.from_dlpack(dst)
    else:
        view = np.asarray(dst)                         # numpy-представление с нулевым копированием

    tile_rows_actual = view.shape[0]                   # не хватает в последнем тайле
    row_start = t * TILE_ROWS                          # начало глобальной оси 0
    row_end = row_start + tile_rows_actual

    # Найти полуоткрытый диапазон индексов файлов, перекрывающийся с [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):
        # Пересечение плитки [row_start, row_end) с файлом [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                # запись в numpy без копирования

manual_task = runtime.create_manual_task(
    load_tile.library,
    load_tile.task_id,
    (num_tasks,),                                      # количество задач в домене запуска == количество тайлов
)
manual_task.add_output(partition)
manual_task.execute()

Оба потребителя проходят через нативные производители PhysicalStore (__dlpack__ для cupy, __array_interface__ для np.asarray) — просмотры локального тайла с нулевым копированием. Сложность бисекции составляет O(log num_shards), а внутренний цикл обычно выполняется 1–2 раза (тайлы перекрываются не более чем на пару файлов).

5. Блокировка и проверка

get_legate_runtime().issue_execution_fence(block=True)

Жёсткие ограничения

1. Все фрагменты должны иметь одинаковый тип данных (dtype) и конечные оси (shape[1:]). Рецепт складывает фрагменты по оси 0; конечные оси места назначения берутся из trailing_shape, которое на этапе обнаружения фиксируется на значении первого файла. Количество строк в каждом фрагменте (shape[0]) может свободно различаться — с этим справляется таблица кумулятивных смещений. В примере любой фрагмент, тип данных или конечная форма которого отличаются от первого, отбрасывается с подробным сообщением об ошибке.

2. Выберите потребитель, соответствующий варианту. cp.from_dlpack отклоняет хранилища, расположенные в SYSMEM; np.asarray без предупреждения возвращает хост-представление хранилища, расположенного в FBMEM, в которое на самом деле нельзя записать данные. Выполняйте диспетчеризацию по ctx.get_variant_kind(), чтобы каждый вариант использовал свой собственный потребитель — см. шаг 4.

3. Представления mmap не всегда являются C-сплошными — оберните каждый фрагмент по файлу с помощью np.ascontiguousarray(arr[file_lo:file_hi]) перед вызовом.set() или записью на месте в numpy.

4. Многоузловая среда: SHARD_DIR должен находиться в общей файловой системе. Каждый рабочий процесс (на каждом ранге) открывает фрагменты по пути; локальные для узла пути /tmp работают только для одноузловых демонстраций.

Варианты

Быстрый путь с однородными фрагментами (одна задача на файл)

Когда все фрагменты уже имеют одинаковые (form, dtype), и у вас имеется количество процессоров, равное num_shards, механизм cum-rows / bisect создаёт лишнюю нагрузку. Установите tile_rows = shard_shape[0] и num_tasks = num_shards; тогда в разбиении будет по одному тайлу на файл, и каждая задача считывает ровно один файл от начала до конца (без бисекции, без внутреннего цикла). Переключение на стороне драйвера выполняется одной строкой:

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 работает в обоих режимах — внутренний цикл просто выполняется ровно один раз на каждую задачу. Нет необходимости в отдельном теле задачи для быстрого пути.

Чрезмерная декомпозиция для лучшей балансировки нагрузки

Значение по умолчанию tile_rows = ceil(total_rows / num_processors) дает по одной плитке на каждый процессор. Чтобы провести избыточную декомпозицию с коэффициентом 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. Описание считывателя: при заданном пути к файлу и полуоткрытом диапазоне строк [file_lo, file_hi) по оси 0 возвращает массив numpy формы (file_hi - file_lo,) + trailing_shape, который можно сделать C-сплошным. Требуются недорогие операции чтения диапазонов/фрагментов — форматы, поддерживающие только «чтение всего файла», не подходят для случая частичного перекрытия (плитка, покрывающая только часть одного файла).

(Встроенные загрузчики с одним вызовом для каждого формата см. в таблице «Зачем нужен этот навык» в начале этого файла.)

На этапе обнаружения (шаг 1) анализируются метаданные каждого формата: .npy / HDF5 / Parquet — все они содержат на диске количество строк в каждом файле и тип данных (dtype). Необработанные двоичные данные — нет; эти сведения извлекаются из размера файла или содержатся в сопутствующем файле.

Распространённые ошибки

cn.asarray(dst) недопустимо в конечной задаче

Внутри тела @task любая операция cupynumeric, затрагивающая верхний уровень среды выполнения — cn.asarray(store), присваивание среза cn_dst[s] = host_np — запускает create_index_space из неправильного контекста, и Legion прерывает работу:

ИСКЛЮЧЕНИЕ ПРИ ИСПОЛЬЗОВАНИИ API LEGION: Недопустимый контекст задачи, переданный вызову среды выполнения
create_index_space

Исправление: используйте капсулу DLPack с помощью сторонней библиотеки (cupy / torch / numpy) внутри листьевых задач. Использование cn.asarray допустимо в драйвере, но не в листьевых задачах. См. файл examples/dlpack/leaf_task_interop.py для обходного решения с использованием torch.

Утверждение (assert) внутри задачи приводит к прерыванию работы среды выполнения

Legate рассматривает негенерированные исключения в @task как нарушение контракта и прерывает выполнение, если задача не была зарегистрирована с помощью throws_exception(). Перед запуском следует выполнить проверку корректности на хосте.

Домен запуска должен соответствовать количеству плиток разбиения

create_manual_task(launch_shape=...) и partition_by_tiling(...) являются независимыми — среда выполнения не обнаруживает несоответствие. Более крупная область запуска → плитки за пределами диапазона; более мелкая → незаписанные плитки. Всегда вычисляйте обе величины из одних и тех же (total_rows, tile_rows) с помощью двух отдельных делений с округлением в большую сторону (если определять размер области запуска напрямую по 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,))