How can I run multiprocessing task defined using @task in Airflow2

Viewed 354

Here is the code:

from datetime import datetime, timedelta
from multiprocessing import Pool

from airflow.operators.python import task
from airflow.operators.python_operator import PythonOperator
from airflow import DAG


def func(a):
    return a + 1


@task
def get_content_body():
    with Pool(2) as p:
        _ = p.map(func, range(10))


def get_content_body2(ti, **context):
    with Pool(2) as p:
        _ = p.map(func, range(10))


default_args = {
    "owner": "Me",
    "depends_on_past": False,
    "start_date": datetime(2021, 1, 1),
    "email_on_failure": False,
    "email_on_retry": False,
    "pool": "default_pool",
    "retries": 1,
    "retry_delay": timedelta(minutes=5),
    "schedule_interval": "@daily",
}

with DAG(f"mydag", catchup=False, default_args=default_args) as dag:
    # task1
    result = get_content_body()

    # task2
    t2 = PythonOperator(
        task_id="my_task2",
        python_callable=get_content_body2,
        provide_context=True,
    )
    t2

As of Airflow 2.x, we can use @task on any function to make it as a task. Problem is that, both tasks defined using @task and using PythonOperator raise an error:

[2021-05-25 04:29:28,768] {taskinstance.py:1455} ERROR - daemonic processes are not allowed to have children
Traceback (most recent call last):
  File "/home/airflow/.local/lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1112, in _run_raw_task
    self._prepare_and_execute_task_with_callbacks(context, task)
  File "/home/airflow/.local/lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1285, in _prepare_and_execute_task_with_callbacks
    result = self._execute_task(context, task_copy)
  File "/home/airflow/.local/lib/python3.8/site-packages/airflow/models/taskinstance.py", line 1315, in _execute_task
    result = task_copy.execute(context=context)
  File "/home/airflow/.local/lib/python3.8/site-packages/airflow/operators/python.py", line 233, in execute
    return_value = self.python_callable(*self.op_args, **self.op_kwargs)
  File "/opt/airflow/dags/repo/dags/dag_custom_plugin_test.py", line 70, in get_content_body
    with Pool(2) as p:
  File "/usr/local/lib/python3.8/multiprocessing/context.py", line 119, in Pool
    return Pool(processes, initializer, initargs, maxtasksperchild,
  File "/usr/local/lib/python3.8/multiprocessing/pool.py", line 212, in __init__
    self._repopulate_pool()
  File "/usr/local/lib/python3.8/multiprocessing/pool.py", line 303, in _repopulate_pool
    return self._repopulate_pool_static(self._ctx, self.Process,
  File "/usr/local/lib/python3.8/multiprocessing/pool.py", line 326, in _repopulate_pool_static
    w.start()
  File "/usr/local/lib/python3.8/multiprocessing/process.py", line 118, in start
    assert not _current_process._config.get('daemon'), \
AssertionError: daemonic processes are not allowed to have children

But the task defined using PythonOperator runs well in Airflow 1.10.x.

How can I make them work in Airflow 2.x?

0 Answers
Related