如何让 airflow dag 等待 VM 在执行下一个任务之前完成其工作

Question

High level description of my workflow

What my vm does -- get data from gcs, process data, save the processed data to gcs

What my dag currently does -- start the vm >> stop the vm >> do the rest of data transformation job.

When I run the above dag, it starts the vm and stops the vm right after that. I would want to make my dag to wait for the vm to finish its job.

Note: kubernates/cloud run is not an option for me.

Answer 1

If you know the output file name which your VM will write to GCS, you can add a sensor between start_vm and stop_vm , which check if the output file is created or not every X seconds, once it's created, the sensor is marked as succeeded and the task stop_vm is started:

gcs_output_sensor = GCSObjectExistenceSensor(
    bucket="bucket_name",
    object="path/to/file"
)
start_the_vm >> gcs_output_sensor >> stop_the_vm >> do_the_rest_of_data_transformation_job

If not, you can instead of running the script which process the data automatically when the VM is up, you can run the VM and use the SSHOperator to run the script, in this case the task will wait the script, and the other task will wait it.