Airflow Composer custom module not found - PythonVirtualenvOperator

Viewed 606

I have a very simple Airflow instance setup in GCP Composer. It has the bucket and everything. I want to set up each dag to run it its own environment with PythonVirtualenvOperator.

The structure in it is as follows:

dags ->
------> code_snippets/
----------> print_name.py - has function called print_my_name() which prints a string into the terminal
------> test_dag.py

test_dag.py:

import datetime
from airflow.operators.python_operator import PythonVirtualenvOperator
from airflow import DAG


def main_func():
    import pandas as pd
    import datetime

    from code_snippets.print_name import print_my_name
    print_my_name()

    df = pd.DataFrame(data={
        'date': [str(datetime.datetime.now().date())]
    })

    print(df)

default_args = {
    'owner': 'test_dag',
    'start_date': datetime.datetime(2020, 7, 3, 5, 1, 00),
    'concurrency': 1,
    'retries': 0
}

dag = DAG('test_dag', description='Test DAGS with environment',
          schedule_interval='0 5 * * *',
          default_args=default_args, catchup=False)

test_the_dag = PythonVirtualenvOperator(
    task_id="test_dag",
    python_callable=main_func,
    python_version='3.8',
    requirements=["DateTime==4.3", "numpy==1.20.2", "pandas==1.2.4", "python-dateutil==2.8.1", "pytz==2021.1",
                  "six==1.15.0", "zope.interface==5.4.0"],
    system_site_packages=False,
    dag=dag,
)

test_the_dag

Everything works until I start importing custom modules - having an init.py does not help, it still gives out the same error, which in my case is:

from code_snippets.print_name import print_my_name\nModuleNotFoundError: No module named \'code_snippets\'

I also have a local instance of Airflow and i experience the same issue. I have tried moving things around or adding the path to the folders to PATH, adding inits in the directories or even changing the import statements, but the error persists as long as I am importing custom modules.

system_site_packages=False or True also has no effect

Is there a fix for that or a way to go around it so I can utilize the custom code I have separated outside of the DAGs?

Airflow Version : 1.10.14+composer Python version for Airflow is set to: 3

1 Answers

The implementation for airflow.operators.python.PythonVirtualenvOperator is such that the python_callable is expected to not reference external names.

Any non-standard library packages used in the callable must be declared as external dependencies in the requirements.txt file.

If you need to use code_snippets, publish it as a package either to pypi or a VCS repository and add it in the list of packages in the requirements kwargs for the PythonVirtualenvOperator.

Related