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, Nov 11 '19 at 16:13
The download_file function it's an API connection that is in another file. — Yaroslav Kolodiy, Nov 11 '19 at 17:41

Matthew Schmill · Answer 1 · 2022-09-24T22:58:26.963

5

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 '22 at 22:58

answered Mar 29 '21 at 17:14

Matthew Schmill

89
1
4

op_kwargs is not valid dictionary in the example you provided. please provide minimum working solution. – Dariusz Krynicki Sep 23 '22 at 11:46
I changed the = to a :. – Matthew Schmill Sep 24 '22 at 23:22
how about python_callable method as he should not use **kwargs but args relating to op_kwargs dict keys. – Dariusz Krynicki Sep 25 '22 at 21:43

score 2 · Answer 2 · edited Apr 23 '21 at 14:46

2

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 '21 at 14:46

raphael

2,762
5
26
55

answered Nov 19 '19 at 16:14

brki

2,672
1
17
12

Thanks for the suggestion, the best way is to use my class "download" as a module to be installed in the virtual operator. – Yaroslav Kolodiy Nov 19 '19 at 16:54
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. – Pelab Dec 20 '20 at 05:24
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 Dec 24 '20 at 09:58
@Pelab if you want to user the PythonVirtualenvOperator verify if you are creating correctly the task [doc](https://github.com/apache/airflow/blob/d4c4492a10f5e269f80067b724aa2c345bf97ee0/airflow/operators/python.py#L361) – Yaroslav Kolodiy Dec 24 '20 at 10:01

How to use PythonVirtualenvOperator in airflow?

2 Answers2

Linked