Vertex.ai自定义训练模型批量预测的Dockerfile配置相关问题

核心逻辑说明 #

批量预测的自定义容器不需要像在线预测那样配置端点、启动web服务，镜像只需要包含可直接执行的预测脚本即可，任务启动后会自动执行脚本，完成后自动退出。 #

1. 输入数据获取方式 #

• Vertex AI 会将你在任务配置阶段指定的输入数据集，自动分片挂载到容器内的/input/data固定路径，你可以直接像读取本地文件一样读取该路径下的所有输入文件。
• 同时平台会自动注入环境变量AIP_BATCH_PREDICTION_INPUT，值为你配置的输入源地址，你也可以在代码中读取该环境变量获取输入路径。
• 若输入源为GCS路径，平台也会自动将对应GCS桶挂载到容器的/gcs/[桶名]/[路径前缀]目录，你也可以直接读取该路径下的文件。

2. 输出结果存放要求 #

• 你只需将预测结果写入容器内的/output/predictions固定路径即可，平台会自动将该目录下的所有文件同步到你任务配置时指定的GCS输出路径。
• 平台同时会注入环境变量AIP_BATCH_PREDICTION_OUTPUT，值为你配置的GCS输出路径，你也可以选择直接通过云SDK将结果写入该路径，更推荐直接写入本地/output/predictions目录，无需额外配置权限。
• 输出格式建议和输入格式对应，常用格式为jsonl（每一行对应一个输入样本的预测结果）或csv格式，平台不会限制输出格式，只要写入指定目录即可。

3. Dockerfile编写规范 #

和在线预测镜像的编写逻辑完全不同，不需要暴露端口、不需要启动web服务，示例如下：

# 按需选择基础运行镜像，这里以Python3.9为例
FROM python:3.9-slim

# 安装预测逻辑需要的依赖
RUN pip install --no-cache-dir pandas scikit-learn

# 将本地预测脚本复制到镜像内
COPY predict.py /app/predict.py

WORKDIR /app

# 直接指定启动预测脚本作为容器启动命令
CMD ["python", "predict.py"]

常见异常原因 #

你遇到的任务运行后抛出异常，大概率是以下两个原因：

• 沿用了在线预测镜像的逻辑，配置了web服务启动命令，没有实际执行预测逻辑，导致任务超时或无输出文件报错
• 预测结果没有写入/output/predictions目录，平台无法采集到输出结果，最终判定任务失败

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