Airflow-sensor-like Tasks

[vsevolodbazhan]

I am currently considering Prefect as an alternative to Airflow. One of the most important features of Airflow for me is Sensors.

Particularly I have used Airflow sensors to wait for files to land on OneDrive/SharePoint or for a reply in a Slack thread.

I have not found any mentions of Sensors in the documentation for Prefect but have found this discussion. API calls are not suited for my use-cases.

I have attempted to recreate Airflow's up_for_reschedule behavior on my own but I am not sure if this is a proper pattern or I can stumble upon some pitfalls.

from datetime import datetime, timedelta
from typing import Optional

import pandas
from prefect import Flow, task
from prefect.engine.results import LocalResult
from prefect.engine.signals import FAIL
from prefect.schedules import IntervalSchedule

IMPORTANT_FILENAME = 'important_file.csv'

schedule = IntervalSchedule(start_date=datetime.utcnow() + timedelta(seconds=1), interval=timedelta(minutes=5))


@task(retry_delay=timedelta(minutes=1), max_retries=3)
def wait_for_data(filename: str) -> pandas.DataFrame:
    try:
        return pandas.read_csv(filepath_or_buffer=filename)
    except FileNotFoundError:
        raise FAIL(message=f'File `{filename}` does not exist.')


@task(result=LocalResult(dir='./sensor', location='{today}'))
def save_data(data: pandas.DataFrame) -> None:
    return data.to_csv()


with Flow('sensors', schedule=schedule) as flow:
    data = wait_for_data(filename=IMPORTANT_FILENAME)
    save_data(data=data)

flow.register(project_name='demo')

[cicdw]

Hi @vsevolodbazhan - great question! You can recreate sensor-like behavior using the following pattern:

import pendulum
from prefect.engine.signals import RETRY


@task
def my_sensor(**kwargs):
    # check some state of the world
    if condition_not_met:
        raise RETRY("Condition not met, retrying in 5 seconds.", start_time=pendulum.now().add(seconds=5))

[vsevolodbazhan]

Thanks for your answer but I have one more question.

Airflow has two sensor modes — poke and reschedule. The first one constantly takes up the worker and the second one frees it up between retries. Which behavior (if either) is more akin to the pattern you have provided?

[vsevolodbazhan]

Another question came to mind. Is there a way to fail a task after a certain number of retries?

[cicdw]

No worries -- the pattern that I have listed above is analogous to poke in that the TaskRunner will actually wait, which takes up a process (Prefect doesn't really have a concept of "worker" in the way that Airflow does - the only cost to waiting / sleeping is a process or job in your environment, and there's no competition from other tasks for this process / job).

If you don't want the process to wait and would rather your Agent resubmit a new process / job, you can use a PAUSE signal instead:

import pendulum
from prefect.engine.signals import PAUSE


@task
def my_sensor(**kwargs):
    # check some state of the world
    if condition_not_met:
        raise PAUSE("Condition not met, retrying in 5 seconds.", start_time=pendulum.now().add(seconds=5))

Although note that for start times < ~1 minute in the future, there's not really a benefit to this over the RETRY signal because it will require your agent to recreate the job / process which could incur overhead depending on the type of agent you use.

[cicdw]

Just saw your follow-up question; yes, for that I would recommend a state handler that checks the retry count and converts to failure after it exceeds a limit, for example:

def retry_handler(task, old, new):
    if new.is_retrying() and new.retry_count > 5:
        raise ValueError("Too many retries")

@task(state_handlers=[retry_handler])
def my_sensor(**kwargs):
    # check some state of the world
    if condition_not_met:
        raise RETRY("Condition not met, retrying in 5 seconds.", start_time=pendulum.now().add(seconds=5))

It doesn't really matter what exception you raise in the state handler, I used a ValueError as an example.

[vsevolodbazhan]

Awesome! Thanks for your answers!