如何在Vertex AI流水线的自定义容器中记录参数？

1. 环境变量传递与记录 #

启动容器时通过-e参数注入训练参数，示例命令：

docker run -e LEARNING_RATE=0.01 -e BATCH_SIZE=32 your-training-image:v1

在训练脚本里读取环境变量并写入日志文件：

import os
import logging

# 配置日志输出到文件
logging.basicConfig(filename='training_params.log', level=logging.INFO)

# 读取参数（带默认值）
lr = float(os.getenv('LEARNING_RATE', 0.001))
batch_size = int(os.getenv('BATCH_SIZE', 16))

# 记录参数
logging.info(f"Training parameters: learning_rate={lr}, batch_size={batch_size}")

容器运行结束后，用docker cp命令把日志文件复制到本地留存即可。

2. 挂载本地配置文件到容器 #

把参数整理到本地的config.yaml或config.json中，启动容器时挂载到容器内指定路径：

docker run -v ./config.yaml:/app/config.yaml your-training-image:v1

训练脚本读取配置文件并记录参数：

import yaml
import logging

logging.basicConfig(filename='training_params.log', level=logging.INFO)

with open('/app/config.yaml', 'r') as f:
    config = yaml.safe_load(f)

logging.info(f"Loaded training parameters: {config}")

这种方式适合参数较多、需要结构化管理的场景。

3. 脚本内直接生成参数记录文件 #

如果参数是脚本内部定义或生成的，直接在训练开始时写入结构化文件：

import json

# 定义或获取训练参数
training_params = {
    "learning_rate": 0.01,
    "batch_size": 32,
    "epochs": 10,
    "model_backbone": "resnet50"
}

# 写入JSON文件
with open('training_params.json', 'w') as f:
    json.dump(training_params, f, indent=4)

训练完成后，通过docker cp提取该文件即可留存参数记录。

4. 轻量参数追踪工具集成 #

如果需要更系统的参数管理，可在容器内集成MLflow（本地运行无需外部服务），脚本中记录参数：

import mlflow

# 设置本地追踪目录
mlflow.set_tracking_uri("./mlruns")
mlflow.start_run()

# 逐一记录参数
mlflow.log_param("learning_rate", 0.01)
mlflow.log_param("batch_size", 32)

# 执行训练逻辑...

mlflow.end_run()

启动容器时把本地目录挂载到容器的./mlruns路径，后续可直接在本地查看所有训练参数和记录。

5. 容器日志直接输出参数 #

在训练脚本中用print或logging直接输出参数，启动容器时将日志重定向到本地文件：

docker run your-training-image:v1 > training_logs.txt 2>&1

脚本中的输出示例：

print(f"Starting training with parameters: lr=0.01, batch_size=32, epochs=10")

这种方式最简便，适合快速验证和基础参数记录需求。

内容的提问来源于stack exchange，提问作者Katie O'Leary