Cache HF models in CI (#410)

Signed-off-by: Xiang Xu <xiangxu@google.com>

Author: xiangxu-google

.buildkite/scripts/run_in_docker.sh

@@ -36,6 +36,20 @@ if [ -z "${MODEL_IMPL_TYPE:-}" ]; then
   MODEL_IMPL_TYPE=flax_nnx
 fi
 
+# Try to cache HF models
+persist_cache_dir="/mnt/disks/persist/models"
+home_cache_dir="$HOME/models"
+
+if ( mkdir -p "$persist_cache_dir" ); then
+  LOCAL_HF_HOME="$persist_cache_dir"
+elif ( mkdir -p "$home_cache_dir" ); then
+  LOCAL_HF_HOME="$home_cache_dir"
+else
+  echo "Error: Failed to create either $persist_cache_dir or $home_cache_dir"
+  exit 1
+fi
+DOCKER_HF_HOME="/tmp/hf_home"
+
 # Prune older images on the host to save space.
 docker system prune -a -f --filter "until=3h"
 
@@ -48,6 +62,8 @@ exec docker run \
   --net host \
   --shm-size=16G \
   --rm \
+  -v "$LOCAL_HF_HOME":"$DOCKER_HF_HOME" \
+  -e HF_HOME="$DOCKER_HF_HOME" \
   -e TPU_BACKEND_TYPE="$TPU_BACKEND_TYPE" \
   -e MODEL_IMPL_TYPE="$MODEL_IMPL_TYPE" \
   -e HF_TOKEN="$HF_TOKEN" \