Vertex AI CreateBatchPredictionJobOperator无法渲染XCom拉取内容的问题

1. 为Vertex AI Operator提供上下文拉取XCom的方法 #

• 尽管CreateBatchPredictionJobOperator不支持provide_context参数，但可以通过两种方式解决XCom拉取问题：
  • 升级Airflow版本：Airflow 2.3.0存在部分Vertex AI Operator的模板渲染缺陷，升级到2.4及以上版本后，model_name字段的模板解析逻辑已修复，可直接使用{{ ti.xcom_pull(key='model_conf')['model_id'] }}语法。
  • 用PythonOperator封装API调用：如果无法升级版本，直接通过PythonOperator调用Vertex AI的批量预测原生API，自由拉取XCom数据，示例代码如下：

from google.cloud import aiplatform
from airflow.operators.python import PythonOperator

def execute_batch_prediction(**context):
    # 从XCom拉取模型ID
    model_info = context['ti'].xcom_pull(key='model_conf')
    full_model_id = model_info['model_id']
    # 初始化Vertex AI环境
    aiplatform.init(project="your-project-id", region="your-region")
    # 执行批量预测
    model = aiplatform.Model(full_model_id)
    batch_job = model.batch_predict(
        job_display_name="custom-batch-pred-job",
        gcs_source="gs://your-bucket/input-data/*.jsonl",
        gcs_destination_prefix="gs://your-bucket/output-results/",
        machine_type="n1-standard-4"
    )
    batch_job.wait()

# 定义任务
batch_pred_task = PythonOperator(
    task_id="run_custom_batch_prediction",
    python_callable=execute_batch_prediction,
    provide_context=True,
    dag=your_dag_object
)

2. 语法错误与文档误解排查 #

• 你使用的{{ ti.xcom_pull(key='model_conf')['model_id'] }}语法本身没有问题，PythonOperator能正常运行也验证了这一点。
• 文档标注model_name为模板化字段但未生效，本质是Airflow 2.3.0的版本bug：该版本的CreateBatchPredictionJobOperator未将model_name加入到template_fields列表中，导致模板引擎不会解析这个字段。
• 额外检查：确保XCom中存储的model_id是完整的Vertex AI资源路径（格式为projects/{project}/locations/{region}/models/{model-id}），如果仅存储了模型ID的后缀部分，需要拼接完整路径后再传入，示例：

{{ "projects/your-project/locations/your-region/models/" + ti.xcom_pull(key='model_conf')['model_id'] }}

内容的提问来源于stack exchange，提问作者no steppin' on snekin'