Как записать большой массив dask (numpy .ndarray) в файл Zarr с использованием графических процессоров?

Question

Я пытаюсь записать большой массив dask (46 ГБ с порциями от 124 до 370 МБ) в файл zarr с помощью dask. Если бы мой массив dask был назван dask_data, тогда работал бы простой dask_data.to_zarr("my_zarr.zarr"). Но насколько я понимаю, это синхронный процесс с привязкой к процессору.

Я хотел бы использовать параллелизм с большей частью работы, выделенной для графического процессора Quadro GV100. Я пытался преобразовать numpy .ndarray в cupy.ndarray через dask_data_cupy = dask_data.map_blocks(cupy.asarray) и записать это в файл zarr, но я получаю:

ValueError: object __array__ method not producing an array (и, честно говоря, я не вижу либо повышение производительности).

Как я могу go использовать графический процессор для распараллеливания записи массива dask в файл zarr?

Спасибо!

Answer 1

Но, насколько я понимаю, это синхронный процесс с привязкой к процессору.

Это, вероятно, неверно, скорее всего, узким местом является устройство хранения. В любом случае каждый кусок записывается в отдельный файл и параллельно между потоками и / или процессами (в зависимости от вашей настройки). В этом весь смысл проекта zarr, что приложение может взаимодействовать с каждым чанком независимо.

Вы можете быть привязанными к процессору, если вы решите использовать различные кодировки сжатия; однако они не обязательно пригодны для работы графического процессора.

Короче говоря, если ваши данные уже не были сгенерированы на графическом процессоре, я был бы удивлен, если передача данных на графический процессор для обработки перед записью в файлы того стоит. Если бы была функция для непосредственного чтения / записи массивов Cupy в Zarr, и вы также обрабатывали на GPU, это было бы иначе - но я не верю, что это так.

Answer 2

Я думаю, вам нужно будет добавить dask_data.map_blocks(cupy.asnumpy) перед вызовом to_zarr.

CuPy пытается убедиться, что пользователь намеревался сделать передачу с устройства на хост (так как это может быть дорого). Таким образом, преднамеренно повышается, когда numpy.asarray вызывается для массива CuPy (как это произошло бы во время этой записи).