Обходим файловый кэш 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). Можно было наверное этот скрипт как то изящнее офрмить, но он сейчас в рабочем состоянии. Может кому-то сэкономит время.

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