JAX может взять код из кеша. Сначала проверьте, кто вправе туда писать

Постоянный кеш JAX помогает повторно использовать скомпилированные программы между запусками. Это удобно для регулярных вычислений: подготовка уже выполнялась, и её не всегда нужно повторять. Но у такого удобства есть условие — каталог нельзя считать безобидной папкой с временными данными.

В руководстве JAX по кешу прямо предупреждают: пользователь, которому разрешена запись в этот каталог, может добиться исполнения произвольного кода вашим процессом. Поэтому первый шаг — определить границу доверия, а уже затем измерять выигрыш.

Подготовьте отдельный каталог

Для небольшого локального задания выберите постоянную папку в своём пользовательском пространстве. Проверьте владельца и права записи самой папки и её родительских каталогов. Если кеш находится на общем томе, выясните, какие другие задания или пользователи могут его менять. Одного знакомого имени каталога для этого недостаточно.

Учебный сценарий: две команды запускают обучение от разных служебных пользователей. Общая папка экономит повторную подготовку, но право записи одной команды становится доверием к её содержимому со стороны другой. Если такое доверие не предусмотрено, начните с отдельных каталогов. Это не сообщение о найденном взломе, а следствие модели безопасности кеша.

Включите кеш до компиляции

from pathlib import Path
import jax

cache_dir = Path.home() / ".cache" / "my-jax-job"
jax.config.update("jax_compilation_cache_dir", str(cache_dir))

Пример задаёт только место хранения; создание папки не доказывает правильность её прав. Настройку нужно выполнить до первой компиляции. Не меняйте разрешения общего каталога вслепую: сначала проверьте, какие другие процессы его используют.

Проверьте пользу двумя запусками

Запустите одно и то же разрешённое задание отдельными процессами, сохранив каталог между ними. Зафиксируйте версии JAX и jaxlib, устройство и форму входных массивов. Сравните подготовку и последующее вычисление отдельно. Быстрый повтор внутри одного процесса ещё не показывает пользу именно дискового кеша.

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

При работе на нескольких машинах отдельно проверьте защиту сетевых служб JAX: права на кеш не заменяют контроль сети и аутентификацию координационного сервиса.

Критерий выбора прост: включайте общий кеш, когда понятны его писатели, сохраняется совместимость заданий и повторная подготовка действительно заметна во времени запуска. Если источник содержимого не контролируется, экономия компиляции не оправдывает расширения доверия. Сам кеш не обещает ускорить каждый шаг уже работающей модели.