How to use PythonVirtualenvOperator in airflow?

Question

Basically I'm working with airflow and developed a task that my download a file from an external source.

t1 = PythonOperator(
        task_id='download',
        python_callable=download,
        provide_context=True,
        dag=dag)

and this airflow is running in a virtual environment (pipenv).

The download function is:

def download(**kwargs):
   folder_id = 'xxxxxx-xxxx-xxxx-xxxxxx'
   file_name = download_file(folder_id)
   return file_name

so basically I'm using Xcons to pass data from one task to another...and using this configurations it's impossible to manage all of dependencies of each DAG...

In the documentation I found this class called "PythonVirtualenvOperator", so to implement that I wrote :

t1 = PythonVirtualenvOperator(
        task_id='download',
        python_callable=download,
        requirements=['requests'],
        python_version='3.8',
        provide_context=True,
        dag=dag
    )

and its giving me the following error:

TypeError: can't pickle module objects

the download_file function it's an API connection that is in another file.

any suggestion how can I manage the environment and have connection between tasks?

What is the download_file function? Can you post your entire DAG for us to help you — kaxil
– kaxil, Commented Nov 11, 2019 at 16:13
The download_file function it's an API connection that is in another file. — Yaroslav Kolodiy
– Yaroslav Kolodiy, Commented Nov 11, 2019 at 17:41

Matthew Schmill · Accepted Answer · 2022-09-24 22:58:26Z

6

The problem is

provide_context=True,

Airflow cannot pickle the context because of all the unserializable stuff in it. You can use templating and op_kwargs to work around this if you only need simple stuff like execution_ts:

t1 = PythonVirtualenvOperator(
        task_id='download',
        python_callable=download,
        provide_context=False,
        op_kwargs={
          execution_date_str: '{{ execution_date }}',
        },
        dag=dag)

Of course, you will need to update the arguments to your callable. I didn't go any deeper than that because it worked for my use case.

edited Sep 24, 2022 at 22:58

answered Mar 29, 2021 at 17:14

Matthew Schmill

991 silver badge4 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

Dariusz Krynicki Over a year ago

op_kwargs is not valid dictionary in the example you provided. please provide minimum working solution.

Matthew Schmill Over a year ago

I changed the = to a :.

Dariusz Krynicki Over a year ago

how about python_callable method as he should not use **kwargs but args relating to op_kwargs dict keys.

raphael · Accepted Answer · 2021-04-23 14:46:11Z

3

From the definition of the PythonVirtualenvOperator:

The function must be defined using def, and not be
part of a class. All imports must happen inside the function
and no variables outside of the scope may be referenced.

I'm guessing that someplace in the chain of code that's being called in your download function, there's a method that's imported from another file using a top-level import. Perhaps moving that import into your download function is enough?

edited Apr 23, 2021 at 14:46

raphael

2,8795 gold badges30 silver badges57 bronze badges

answered Nov 19, 2019 at 16:14

brki

2,7901 gold badge19 silver badges12 bronze badges

4 Comments

Yaroslav Kolodiy Over a year ago

Thanks for the suggestion, the best way is to use my class "download" as a module to be installed in the virtual operator.

Pelab Over a year ago

Hi @YaroslavKolodiy, I am facing an issue with PythonVirtualenvOperator where the task is not using mention packages and mentioned python version insdie the task. It's looks like operator just skipping the option and using default python version n libraries present in the server where airflow is installed. Can you provide some suggestion on this to me.

Yaroslav Kolodiy Over a year ago

Hi @Pelab I resolved my problem using containerization, I decided to migrate all of my tasks to DockerOperator that have the scripts and dependencies already installed.

Yaroslav Kolodiy Over a year ago

@Pelab if you want to user the PythonVirtualenvOperator verify if you are creating correctly the task doc

Collectives™ on Stack Overflow

How to use PythonVirtualenvOperator in airflow?

2 Answers 2

3 Comments

4 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

4 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related