Unable to store Airflow task objects to a dictionary

Viewed 229

I'm trying to store Airflow task objects in dictionary as values like the following

operator_dict = {}

operator_dict['file-sensor1'] = FileSensor(task_id=file-sensor-1, fs_conn_id='fs_source', filepath=file, poke_interval=30, timeout=300,dag=dag) 

Printing operator_dict dictionary looks like the following

{'filetrigger_1': <Task(FileSensor): filetrigger_1>, 'filetrigger_2': <Task(FileSensor): filetrigger_2>, 'datamorphjob_1': <Task(DatabricksRunNowOperator): datamorphjob_1>})

I see no error/exception raised in the above step so I assume this is allowed and was able to set operators to a dictionary value and then store that entire dictionary to airflow Variable using

Variable.set("my_list", operator_dict)

But when I retrieve the same using Variable.get("my_dict") I assume it is not able to deserialize the dictionary data properly just like it does with the json list data.

For list anyhow I use something like ast.literal_eval(Variable.get("workflow")) and it works but for dictionary, I tried using eval(Variable.get("my_dict")) to identify the issue and it throws the following exception.

enter image description here

Eventually, I want to do my_dict['filetrigger_1'] and retrieve the <Task...> object associated with that and execute it using either of upstream / downstream

I found it difficult to identify the real problem here

1 Answers

Airflow variables have their values converted to strings before being persisted (you can see it happen in source code of Variable.set: https://airflow.apache.org/docs/apache-airflow/stable/_modules/airflow/models/variable.html#Variable.set). As you've seen, the string representation of a Task object can't be deserialized back into the original Task object.

The workaround for this is to pass the arguments for the task instead of passing the task itself. I.e.

operator_dict['files_sensor1'] = {'task_id': file-sensor-1, ...}}

Then, downstream, you would construct the FileSensor via something like

sensor = FileSensor(**operator_dict['files_sensor1'])

But stepping back a bit: this seems like a weird pattern that could cause a lot of issues. What if you try to rerun a task from a previous DAG and it gets the variables you set from a future DAG? Or vice-versa? I'd strongly recommend that you consider using XCOM variables or finding a way to construct the tasks entirely within their own associated DAG.

Related