Обходим файловый кэш Google Drive в Colab

Про colab знают, наверное, все. Этот инструмент позволяет независимым исследователям использовать облачную инфраструктуру с GPU и TPU бесплатно или почти бесплатно.

Как всегда, проблемы возникают на больших данных. Если ваш датасэт лежит в google drive (он же Диск), то вы можете обращаться к нему напрямую из colab. Однако, если файл велик, например, 70+ GiB, то процесс обучения будет существенно медленнее, чем если бы этот же файл лежал в локальном хранилище, которое выделяется при создании инстанса.

Выход — скопировать файл с Диска в локальное хранилище (обучение станет быстрее в несколько раз!). Но дело в том, что colab и вся инфраструктура очень умная, файлы с Диска кэшируются каким то неуправляемым вами алгоритмом. И если у вашего инстанса, допустим, доступно ~120 GiB, то 70 GiB с Диска вы не скопируете, у вас закончится свободное место как раз из-за системы кэширования. То есть, команда cp не отработает корректно. И rsync то же. И tar. Кэширование работает на уровне драйвера. По сути файл копируется в локальное хранилище дважды. Шах и мат!

Так что вот вам небольшой костылёк:

sh = """ #!/bin/bash  BIGFILE=/content/drive/MyDrive/DATA/combined/2021-12-29-combined.h5 OUTFILE=/content/2021-12-29-combined.h5 PARTSIZE=$((1*1024*1024*1024)) # 1 GiB FILESIZE=$(stat -c%s "$BIGFILE") echo "Size of $BIGFILE = $FILESIZE bytes."  CHUNKS=$((FILESIZE / PARTSIZE)) if (( FILESIZE % PARTSIZE == 0 )) then     echo "CHUNKS = $CHUNKS" else     ((CHUNKS = CHUNKS + 1))     echo "CHUNKS = $CHUNKS" fi  for ((i=0;i<$CHUNKS;i++)); do echo "Transfer CHUNK $i out of $CHUNKS"   dd if=$BIGFILE bs=$PARTSIZE skip=$i count=1 iflag=sync oflag=sync >> $OUTFILE   sleep 10   sync done """ with open('script.sh', 'w') as file:   file.write(sh)  !bash script.sh

Можно прямо в ячейку вставлять! Идея в том, что бы копировать файл по частям, с небольшой задержкой. Тогда кэш переодически сбрасывается и всё нормально. Небольшой оверхэд всё равно есть (порядка 10-15 GiB). Можно было наверное этот скрипт как то изящнее офрмить, но он сейчас в рабочем состоянии. Может кому-то сэкономит время.

С Новым годом!