PythonVirtualenvOperator fails with TypeError: cannot pickle 'module' object

Question

The following happens on Cloud Composer 2.1.xx

I am trying to use the PythonVirtualenvOperator with templated parameters. Unfortunately the operator fails with the following error :

TypeError: cannot pickle 'module' object

Here's the code of my dag :

with models.DAG(
    'name',
    schedule_interval=None,
    start_date=datetime.datetime(2018, 1, 1),
    catchup=False) as dag:

t1 = PythonVirtualenvOperator(
    task_id='download',
    python_callable=update_files,
    requirements=['google-cloud-firestore==2.2.0'],
    python_version='3.8',
    op_kwargs={
        "inputPaths": '{{ dag_run.conf["inputPaths"] }}'
    }
)

And the code for the python function looks like this :

def update_files(**kwargs):
    from google.cloud import firestore
    import datetime
    paths = kwargs['inputPaths']
    .....

I tried to use the parameter use_dill = True and the suggested answer on How to use PythonVirtualenvOperator in airflow? but it doesn't make anything better.

Answer 1

From airflow 2.x version, we will not be able to get op_kwargs dict in virtualEnv python_callable method. Instead, we will be able to get only value in the callable In your case, you have change only your callable as

def update_files(inputPaths): // def callable(<DICT_KEY>):
from google.cloud import firestore
import datetime
print(inputPaths)