GCP启动脚本执行报错：bufio.Scanner: token too long 求助

问题描述 #

在GCP上创建实例并通过启动脚本运行机器学习训练任务，实例创建命令如下：

gcloud compute instances create XXXXX \
  --project YYYYY \
  --machine-type='a2-highgpu-1g' \
  --zone='us-central1-a' \
  --image-project='AAAAA' \
  --image-family='BBBBBB' \
  --boot-disk-size=50GB \
  --accelerator type=CCCCCCC,count=1 \
  --metadata "DDDDDDDD" \
  --maintenance-policy TERMINATE --restart-on-failure \
  --scopes https://www.googleapis.com/auth/cloud-platform \
  --metadata-from-file startup-script=start-script.sh

脚本前期安装步骤日志正常，但执行到末尾的Dreambooth训练命令时，约10分钟后报错：

google_metadata_script_runner抛出“error while communicating with "startup-script" script: bufio.Scanner: token too long”，同时标注startup-script退出状态为0

定位到的训练命令如下（手动执行无问题，脚本为unix格式）：

. /root/work/venv_diffusers_sd_2/bin/accelerate launch /root/work/diffusers_sd_v2/examples/dreambooth/train_dreambooth.py \
 --gradient_accumulation_steps=1 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-base" \
 --pretrained_vae_name_or_path "stabilityai/sd-vae-ft-mse" --output_dir=/root/work/train_1/model_out/ --with_prior_preservation \
 --prior_loss_weight=1.0 --resolution=512 --train_batch_size=1 --learning_rate=2e-6 \
 --lr_scheduler="constant" --lr_warmup_steps=0 --num_class_images=200 --max_train_steps=2000 \
 --concepts_list="/root/work/train_1/concepts_list.json" --train_text_encoder --revision="fp16" --mixed_precision="fp16"

原因分析 #

GCP的google_metadata_script_runner在捕获启动脚本的输出时，使用的bufio.Scanner默认最大允许的单行token长度为64KB。训练过程中，模型加载、日志输出可能产生超长单行内容，超过这个限制就会触发报错。而手动执行时输出直接交给shell处理，没有这个长度限制，所以不会报错。

解决方案 #

方案1：重定向训练命令输出到文件 #

修改训练命令，将标准输出和错误输出全部重定向到日志文件，避免google_metadata_script_runner捕获到超长行：

. /root/work/venv_diffusers_sd_2/bin/accelerate launch /root/work/diffusers_sd_v2/examples/dreambooth/train_dreambooth.py \
 --gradient_accumulation_steps=1 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-base" \
 --pretrained_vae_name_or_path "stabilityai/sd-vae-ft-mse" --output_dir=/root/work/train_1/model_out/ --with_prior_preservation \
 --prior_loss_weight=1.0 --resolution=512 --train_batch_size=1 --learning_rate=2e-6 \
 --lr_scheduler="constant" --lr_warmup_steps=0 --num_class_images=200 --max_train_steps=2000 \
 --concepts_list="/root/work/train_1/concepts_list.json" --train_text_encoder --revision="fp16" --mixed_precision="fp16" > /root/work/train_logs.txt 2>&1

后续可通过cat /root/work/train_logs.txt查看训练日志。

方案2：后台执行训练命令 #

使用nohup让训练命令在后台运行，同时脱离启动脚本的shell环境，避免输出被扫描：

nohup . /root/work/venv_diffusers_sd_2/bin/accelerate launch /root/work/diffusers_sd_v2/examples/dreambooth/train_dreambooth.py \
 --gradient_accumulation_steps=1 --pretrained_model_name_or_path="stabilityai/stable-diffusion-2-base" \
 --pretrained_vae_name_or_path "stabilityai/sd-vae-ft-mse" --output_dir=/root/work/train_1/model_out/ --with_prior_preservation \
 --prior_loss_weight=1.0 --resolution=512 --train_batch_size=1 --learning_rate=2e-6 \
 --lr_scheduler="constant" --lr_warmup_steps=0 --num_class_images=200 --max_train_steps=2000 \
 --concepts_list="/root/work/train_1/concepts_list.json" --train_text_encoder --revision="fp16" --mixed_precision="fp16" > /root/work/train_logs.txt 2>&1 &

此方法保证启动脚本结束后训练任务仍继续运行，适合长时间训练场景。

方案3：调整扫描缓冲区大小（不推荐） #

修改google_metadata_script_runner的代码或配置，增大bufio.Scanner的最大token长度。但此操作需要修改系统文件，实例重启后会失效，且风险较高，一般不建议采用。

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