如何在GCP AI Platform训练作业中挂载ram-disk/tmpfs

gcloud ai-platform jobs submit training 确实没有提供类似Compute Engine容器挂载tmpfs的原生参数，但完全可以通过容器内自主挂载的方式实现ram-disk，效果和平台侧配置的tmpfs完全一致。
AI Platform训练任务的自定义容器运行时默认持有CAP_SYS_ADMIN权限，没有封禁mount系统调用，不需要额外申请权限就能完成tmpfs挂载。

方案1：修改容器入口点（推荐，一劳永逸） #

直接在容器的启动流程里加入tmpfs挂载逻辑，不需要每次提交任务都写额外命令：

1. 在项目里新建启动脚本entrypoint.sh，内容如下：

   #!/bin/bash
   set -e
   # 提前创建挂载目录，目录不存在会挂载失败
   mkdir -p /cache
   # 挂载tmpfs，size参数根据你选的训练机型内存调整，建议预留30%左右内存给训练进程和系统开销
   # 例：机型总内存32G就设20G，总内存128G就设90G，避免触发OOM
   mount -t tmpfs -o size=20G tmpfs /cache
   # 执行原有训练启动命令
   exec "$@"
   

2. 修改你的训练镜像Dockerfile，把脚本加进去设为入口点：

   # 把启动脚本复制到镜像内
   COPY entrypoint.sh /opt/entrypoint.sh
   RUN chmod +x /opt/entrypoint.sh
   # 设为入口点
   ENTRYPOINT ["/opt/entrypoint.sh"]
   # 保留你原来的训练启动命令即可，比如
   CMD ["python", "train.py"]
   

3. 重新用Cloud Build构建镜像推到GCR，提交训练任务时正常用这个镜像就行，容器启动后会自动把/cache挂载为内存盘。

方案2：提交任务时直接写启动命令（不用改镜像） #

如果不想重新构建镜像，可以在提交训练任务时，把启动命令替换为先挂载再跑训练的逻辑，示例提交命令：

gcloud ai-platform jobs submit training tmpfs_test_job \
  --region us-central1 \
  --scale-tier CUSTOM \
  --master-machine-type n1-standard-8 \
  --master-image-uri gcr.io/your-project-id/your-train-image:v1 \
  -- \
  bash -c "mkdir -p /cache && mount -t tmpfs -o size=20G tmpfs /cache && python train.py"

常见场景补充 #

如果是PyTorch/TensorFlow多进程数据加载需要调大共享内存/dev/shm，不需要新建目录，直接在启动时重挂载即可，命令如下：

# 把/dev/shm调整为16G
mount -o remount,size=16G /dev/shm

注意：tmpfs占用的是实例的物理内存，不要把size设得超过实例可用内存上限，否则会触发系统OOM直接杀掉训练进程。

内容的提问来源于stack exchange，提问作者Stiefel