在Colab运行Dataflow任务时遇到flexrs_goal参数异常错误

问题现象 #

在Colab运行Dataflow任务时，Worker抛出参数错误：

sdk_worker_main.py: error: argument --flexrs_goal: invalid choice: '/root/.local/share/jupyter/runtime/kernel-1dbd101c-a79e-432e-89b3-5ba68df104d7.json' (choose from 'COST_OPTIMIZED', 'SPEED_OPTIMIZED')

未手动设置flexrs_goal参数，手动指定该参数、回滚apache-beam版本均无效。相关配置与代码如下：

Pipeline配置 #

beam_options = PipelineOptions(
    runner='DataflowRunner',
    project=...,
    job_name=...,
    temp_location=...,
    subnetwork='regions/us-west1/subnetworks/default',
    region='us-west1'
)

Pipeline代码 #

with beam.Pipeline(options=beam_options) as pipeline:
  (pipeline 
    | beam.io.ReadFromBigQuery(
        query=f'SELECT column FROM {BQ_TABLE} LIMIT 100')
    | beam.Map(print))

问题本质是Colab的Jupyter环境会自动将内核相关参数传递给Dataflow Worker，导致flexrs_goal被错误赋值为内核JSON文件路径。

解决方案 #

方案1：显式指定参数并阻断环境参数传递 #

在PipelineOptions中添加合法的flexrs_goal值，同时关闭save_main_session以阻止Jupyter传递额外参数：

beam_options = PipelineOptions(
    runner='DataflowRunner',
    project=...,
    job_name=...,
    temp_location=...,
    subnetwork='regions/us-west1/subnetworks/default',
    region='us-west1',
    flexrs_goal='COST_OPTIMIZED',  # 可选值：COST_OPTIMIZED/SPEED_OPTIMIZED
    save_main_session=False
)

方案2：清理Colab环境中的干扰变量 #

在笔记本开头添加代码，移除可能传递给Worker的Jupyter相关环境变量：

import os
if 'PYTHONPATH' in os.environ:
    os.environ['PYTHONPATH'] = ':'.join([p for p in os.environ['PYTHONPATH'].split(':') if not '/root/.local/share/jupyter' in p])

方案3：重启内核后切换运行环境测试 #

1. 先将runner改为DirectRunner运行Pipeline，验证逻辑无错误
2. 通过「Runtime -> Restart runtime」重启Colab内核
3. 改回DataflowRunner重新提交任务，避免残留的Jupyter参数干扰

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