如何让Airflow DAG等待VM完成任务后再执行后续任务

针对你的场景，这里提供几个无需Kubernetes/Cloud Run的可行方案，让DAG等待VM完成数据处理后再执行后续步骤：

方案1：GCS标记文件+Airflow传感器 #

这是最直观的实现方式，让VM完成任务后在GCS写入一个标记文件，Airflow通过传感器等待该文件出现：

1. 修改VM启动脚本：
   在VM的数据处理脚本末尾添加代码，往指定GCS路径上传一个完成标记文件，比如：

   # 数据处理完成后执行
   gsutil touch gs://your-bucket-name/vm-task-signal/task_completed.done
   

   确保VM的服务账号拥有GCS写入权限。

2. 调整Airflow DAG流程：
   将原有的「启动VM >> 停止VM」改为如下依赖链：

   from airflow import DAG
   from airflow.providers.google.cloud.operators.compute import ComputeEngineStartInstanceOperator, ComputeEngineStopInstanceOperator
   from airflow.providers.google.cloud.sensors.gcs import GoogleCloudStorageObjectSensor
   from datetime import datetime
   
   default_args = {
       'start_date': datetime(2024, 1, 1)
   }
   
   with DAG('vm_data_processing', default_args=default_args, schedule_interval='@daily') as dag:
       start_vm = ComputeEngineStartInstanceOperator(
           task_id='start_vm',
           project_id='your-gcp-project',
           zone='us-central1-a',
           resource_id='your-vm-instance'
       )
   
       wait_for_completion = GoogleCloudStorageObjectSensor(
           task_id='wait_for_vm_completion',
           bucket='your-bucket-name',
           object='vm-task-signal/task_completed.done',
           mode='poke',
           poke_interval=60
       )
   
       stop_vm = ComputeEngineStopInstanceOperator(
           task_id='stop_vm',
           project_id='your-gcp-project',
           zone='us-central1-a',
           resource_id='your-vm-instance'
       )
   
       # 剩余的数据转换任务
       data_transform = ...  # 你的转换任务Operator
   
       # 设置依赖
       start_vm >> wait_for_completion >> stop_vm >> data_transform
   

   每次运行前可以先清理旧的标记文件，避免重复触发。

方案2：VM元数据状态监控 #

让VM完成任务后更新自身的自定义元数据，Airflow通过轮询元数据状态来判断任务是否完成：

1. VM内更新元数据：
   处理完成后执行gcloud命令更新元数据：

   gcloud compute instances add-metadata your-vm-instance \
       --zone us-central1-a \
       --metadata task_status=completed
   

2. Airflow中轮询元数据：
   用PythonOperator调用GCP Compute API检查元数据状态：

   from google.cloud import compute_v1
   from datetime import timedelta
   from airflow.operators.python import PythonOperator
   
   def check_vm_task_status(**context):
       instance_client = compute_v1.InstancesClient()
       instance = instance_client.get(
           project='your-gcp-project',
           zone='us-central1-a',
           instance='your-vm-instance'
       )
       # 获取自定义元数据
       task_status = None
       for item in instance.metadata.items:
           if item.key == 'task_status':
               task_status = item.value
               break
       if task_status != 'completed':
           raise ValueError("VM task not completed yet")
   
   # 在DAG中添加任务
   check_status = PythonOperator(
       task_id='check_vm_task_status',
       python_callable=check_vm_task_status,
       provide_context=True,
       retries=10,
       retry_delay=timedelta(minutes=1)
   )
   
   # 依赖链：start_vm >> check_status >> stop_vm >> data_transform
   

方案3：Cloud Logging日志触发 #

利用VM输出的特定日志作为完成信号，Airflow通过日志传感器等待该信号：

1. VM输出完成日志：
   在处理脚本末尾添加日志输出：

   echo "[VM_TASK_DONE] Data processing and upload to GCS completed successfully"
   

2. Airflow配置日志传感器：
   使用CloudLoggingSensor过滤该日志：

   from airflow.providers.google.cloud.sensors.logging import CloudLoggingSensor
   
   wait_for_log = CloudLoggingSensor(
       task_id='wait_for_vm_log',
       project_id='your-gcp-project',
       filter_='textPayload:"[VM_TASK_DONE]"',
       poke_interval=60,
       timeout=3600
   )
   
   # 依赖链：start_vm >> wait_for_log >> stop_vm >> data_transform
   

注意事项 #

• 确保Airflow的服务账号拥有对应GCP资源的访问权限（GCS读取、Compute实例元数据读取、Logging读取等）
• 根据VM任务的实际时长，调整传感器的轮询间隔和超时时间
• 每次DAG运行前，建议重置VM的标记文件/元数据状态，避免上次运行的残留数据影响本次判断

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