Run tasks in an isolated environment in Apache Airflow

1
# from airflow.operators.python import ExternalPythonOperator
2
# import os
3
4
def my_isolated_function():
5
    print(f"The python version in the virtual env is: {sys.version}")
6
    print(f"The pandas version in the virtual env is: {pd.__version__}")
7
8
my_isolated_task = ExternalPythonOperator(
9
    task_id="my_isolated_task",
10
    python_callable=my_isolated_function,
11
    python=os.environ["ASTRO_PYENV_epo_pyenv"]
12
)

1
"""
2
## Toy example of using the @task.external_python decorator
3
4
The @task.external_python decorator is used to run any Python code in an existing isolated Python environment.
5
"""
6
7
from airflow.decorators import dag, task
8
import pandas as pd
9
import sys
10
import os
11
12
13
@dag(
14
    start_date=None,
15
    schedule=None,
16
    doc_md=__doc__,
17
    description="@task.external_python",
18
    default_args={
19
        "owner": "airflow",
20
        "retries": 0,
21
    },
22
    tags=["@task.external_python"],
23
)
24
def external_python_decorator_dag():
25
26
    @task
27
    def upstream_task():
28
        print(f"The python version in the upstream task is: {sys.version}")
29
        print(f"The pandas version in the upstream task is: {pd.__version__}")
30
        return {"num": 1, "word": "hello"}
31
32
    @task.external_python(python=os.environ["ASTRO_PYENV_epo_pyenv"])
33
    def my_isolated_task(upstream_task_output: dict):
34
        """
35
        This function runs in an isolated environment.
36
        Args:
37
            upstream_task_output (dict): contains a number and a word.
38
        Returns:
39
            pd.DataFrame: A dictionary containing the transformed inputs.
40
        """
41
        import pandas as pd
42
        import sys
43
44
        print(f"The python version in the virtual env is: {sys.version}")
45
        print(f"The pandas version in the virtual env is: {pd.__version__}")
46
47
        num = upstream_task_output["num"]
48
        word = upstream_task_output["word"]
49
50
        num_plus_one = num + 1
51
        word_plus_exclamation = word + "!"
52
53
        df = pd.DataFrame(
54
            {
55
                "num_plus_one": [num_plus_one],
56
                "word_plus_exclamation": [word_plus_exclamation],
57
            },
58
        )
59
60
        return df
61
62
    @task
63
    def downstream_task(arg):
64
        print(f"The python version in the downstream task is: {sys.version}")
65
        print(f"The pandas version in the downstream task is: {pd.__version__}")
66
        return arg
67
68
    downstream_task(my_isolated_task(upstream_task()))
69
70
71
external_python_decorator_dag()

1
"""
2
## Toy example of using the ExternalPythonOperator
3
4
The ExternalPythonOperator is used to run any Python code in an existing isolated Python environment.
5
"""
6
7
from airflow.decorators import dag, task
8
from airflow.models.baseoperator import chain
9
from airflow.operators.python import ExternalPythonOperator
10
import pandas as pd
11
import sys
12
import os
13
14
15
def my_isolated_function(num: int, word: str) -> dict:
16
    """
17
    This function will be passed to the ExternalPythonOperator to
18
    run in an isolated environment.
19
    Args:
20
        num (int): An integer to be incremented by 1.
21
        word (str): A string to have an exclamation mark added to it.
22
    Returns:
23
        pd.DataFrame: A dictionary containing the transformed inputs.
24
    """
25
    import pandas as pd
26
    import sys
27
28
    print(f"The python version in the virtual env is: {sys.version}")
29
    print(f"The pandas version in the virtual env is: {pd.__version__}")
30
31
    num_plus_one = num + 1
32
    word_plus_exclamation = word + "!"
33
34
    df = pd.DataFrame(
35
        {
36
            "num_plus_one": [num_plus_one],
37
            "word_plus_exclamation": [word_plus_exclamation],
38
        },
39
    )
40
41
    return df
42
43
44
@dag(
45
    start_date=None,
46
    schedule=None,
47
    doc_md=__doc__,
48
    description="ExternalPythonOperator",
49
    render_template_as_native_obj=True,
50
    default_args={
51
        "owner": "airflow",
52
        "retries": 0,
53
    },
54
    tags=["ExternalPythonOperator"],
55
)
56
def external_python_operator_dag():
57
58
    @task
59
    def upstream_task():
60
        print(f"The python version in the upstream task is: {sys.version}")
61
        print(f"The pandas version in the upstream task is: {pd.__version__}")
62
        return {"num": 1, "word": "hello"}
63
64
    my_isolated_task = ExternalPythonOperator(
65
        task_id="my_isolated_task",
66
        python_callable=my_isolated_function,
67
        python=os.environ["ASTRO_PYENV_epo_pyenv"],
68
        op_kwargs={
69
            # note that render_template_as_native_obj=True in the DAG definition
70
            # to render num as an integer
71
            "num": "{{ ti.xcom_pull(task_ids='upstream_task')['num']}}",
72
            "word": "{{ ti.xcom_pull(task_ids='upstream_task')['word']}}",
73
        },
74
    )
75
76
    @task
77
    def downstream_task(arg):
78
        print(f"The python version in the downstream task is: {sys.version}")
79
        print(f"The pandas version in the downstream task is: {pd.__version__}")
80
        return arg
81
82
    chain(upstream_task(), my_isolated_task, downstream_task(my_isolated_task.output))
83
84
85
external_python_operator_dag()

1
# from airflow.operators.python import PythonVirtualenvOperator
2
3
def my_isolated_function():
4
    print(f"The pandas version in the virtual env is: {pd.__version__}")
5
    # your code to run in the isolated environment
6
7
my_isolated_task = PythonVirtualenvOperator(
8
    task_id="my_isolated_task",
9
    python_callable=my_isolated_function,
10
    requirements=[
11
        "pandas==1.5.1",
12
    ]  # add your requirements to the list
13
)

1
"""
2
## Toy example of using the @task.virtualenv decorator
3
4
The @task.virtualenv decorator is used to run any Python code in a new isolated Python environment.
5
"""
6
7
from airflow.decorators import dag, task
8
import pandas as pd
9
10
11
@dag(
12
    start_date=None,
13
    schedule=None,
14
    doc_md=__doc__,
15
    description="@task.virtualenv",
16
    default_args={
17
        "owner": "airflow",
18
        "retries": 0,
19
    },
20
    tags=["@task.virtualenv"],
21
)
22
def virtualenv_decorator_dag():
23
24
    @task
25
    def upstream_task():
26
        print(f"The pandas version in the upstream task is: {pd.__version__}")
27
        return {"num": 1, "word": "hello"}
28
29
    @task.virtualenv(requirements=["pandas==1.5.1"])
30
    def my_isolated_task(upstream_task_output: dict):
31
        """
32
        This function runs in an isolated environment.
33
        Args:
34
            upstream_task_output (dict): contains a number and a word.
35
        Returns:
36
            pd.DataFrame: A dictionary containing the transformed inputs.
37
        """
38
        import pandas as pd
39
40
        print(f"The pandas version in the virtual env is: {pd.__version__}")
41
42
        num = upstream_task_output["num"]
43
        word = upstream_task_output["word"]
44
45
        num_plus_one = num + 1
46
        word_plus_exclamation = word + "!"
47
48
        df = pd.DataFrame(
49
            {
50
                "num_plus_one": [num_plus_one],
51
                "word_plus_exclamation": [word_plus_exclamation],
52
            },
53
        )
54
55
        return df
56
57
    @task
58
    def downstream_task(arg):
59
        print(f"The pandas version in the downstream task is: {pd.__version__}")
60
        return arg
61
62
    downstream_task(my_isolated_task(upstream_task_output=upstream_task()))
63
64
65
virtualenv_decorator_dag()

1
"""
2
## Toy example of using the PythonVirtualenvOperator
3
4
The PythonVirtualenvOperator is used to run any Python code in a new isolated Python environment.
5
"""
6
7
from airflow.decorators import dag, task
8
from airflow.models.baseoperator import chain
9
from airflow.operators.python import PythonVirtualenvOperator
10
import pandas as pd
11
import sys
12
13
14
def my_isolated_function(num: int, word: str) -> dict:
15
    """
16
    This function will be passed to the PythonVirtualenvOperator to
17
    run in an isolated environment.
18
    Args:
19
        num (int): An integer to be incremented by 1.
20
        word (str): A string to have an exclamation mark added to it.
21
    Returns:
22
        pd.DataFrame: A dictionary containing the transformed inputs.
23
    """
24
    import pandas as pd
25
26
    print(f"The pandas version in the virtual env is: {pd.__version__}")
27
28
    num_plus_one = num + 1
29
    word_plus_exclamation = word + "!"
30
31
    df = pd.DataFrame(
32
        {
33
            "num_plus_one": [num_plus_one],
34
            "word_plus_exclamation": [word_plus_exclamation],
35
        },
36
    )
37
38
    return df
39
40
41
@dag(
42
    start_date=None,
43
    schedule=None,
44
    doc_md=__doc__,
45
    description="PythonVirtualenvOperator",
46
    render_template_as_native_obj=True,
47
    default_args={
48
        "owner": "airflow",
49
        "retries": 0,
50
    },
51
    tags=["PythonVirtualenvOperator"],
52
)
53
def python_virtualenv_operator_dag():
54
55
    @task
56
    def upstream_task():
57
        print(f"The python version in the upstream task is: {sys.version}")
58
        print(f"The pandas version in the upstream task is: {pd.__version__}")
59
        return {"num": 1, "word": "hello"}
60
61
    my_isolated_task = PythonVirtualenvOperator(
62
        task_id="my_isolated_task",
63
        python_callable=my_isolated_function,
64
        requirements=["pandas==1.5.1"],
65
        op_kwargs={
66
            # note that render_template_as_native_obj=True in the DAG definition
67
            # to render num as an integer
68
            "num": "{{ ti.xcom_pull(task_ids='upstream_task')['num']}}",
69
            "word": "{{ ti.xcom_pull(task_ids='upstream_task')['word']}}",
70
        },
71
    )
72
73
    @task
74
    def downstream_task(arg):
75
        print(f"The python version in the downstream task is: {sys.version}")
76
        print(f"The pandas version in the downstream task is: {pd.__version__}")
77
        return arg
78
79
    chain(upstream_task(), my_isolated_task, downstream_task(my_isolated_task.output))
80
81
82
python_virtualenv_operator_dag()

1
#from airflow.operators.python import PythonVirtualenvOperator
2
3
def my_isolated_function():
4
    print(f"The python version in the virtual env is: {sys.version}")
5
    print(f"The pandas version in the virtual env is: {pd.__version__}")
6
    # your code to run in the isolated environment
7
8
my_isolated_task = PythonVirtualenvOperator(
9
    task_id="my_isolated_task",
10
    python_callable=my_isolated_function,
11
    requirements=["pandas==1.5.1"],
12
    python_version="3.10",  # specify the Python version
13
)

1
# from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
2
# from airflow.configuration import conf
3
4
# if you are running Airflow on Kubernetes, you can get 
5
# the current namespace from the Airflow conf
6
namespace = conf.get("kubernetes", "NAMESPACE")
7
8
my_isolated_task = KubernetesPodOperator(
9
    task_id="my_isolated_task",
10
    namespace=namespace,
11
    # your Docker image contains the scripts to run in the isolated environment  
12
    image="<YOUR IMAGE>",
13
    name="<YOUR POD NAME>",
14
    in_cluster=True,
15
    is_delete_operator_pod=True,
16
    get_logs=True,
17
)

1
# from airflow.decorators import task
2
# import os
3
4
@task.branch_external_python(python=os.environ["ASTRO_PYENV_epo_pyenv"])
5
def my_isolated_task():
6
    print(f"The pandas version in the virtual env is: {pd.__version__}")
7
8
    num = random.randint(0, 100)
9
10
    if num > 50:
11
        # return the task_id of the downstream task that should be executed
12
        return "downstream_task_a"  
13
    else:
14
        return "downstream_task_b"

1
# from airflow.operators.python import BranchExternalPythonOperator
2
# import os
3
4
def my_isolated_function():
5
    print(f"The pandas version in the virtual env is: {pd.__version__}")
6
7
    num = random.randint(0, 100)
8
9
    if num > 50:
10
        # return the task_id of the downstream task that should be executed
11
        return "downstream_task_a"  
12
    else:
13
        return "downstream_task_b"
14
15
my_isolated_task = BranchExternalPythonOperator(
16
    task_id="my_isolated_task",
17
    python_callable=my_isolated_function,
18
    python=os.environ["ASTRO_PYENV_epo_pyenv"]
19
)

1
# from airflow.decorators import task
2
3
@task.branch_virtualenv(requirements=["pandas==1.5.3"])
4
def my_isolated_task():
5
    print(f"The pandas version in the virtual env is: {pd.__version__}")
6
7
    num = random.randint(0, 100)
8
9
    if num > 50:
10
        # return the task_id of the downstream task that should be executed
11
        return "downstream_task_a"
12
    else:
13
        return "downstream_task_b"

1
# from airflow.operators.python import BranchPythonVirtualenvOperator
2
3
def my_isolated_function():
4
    print(f"The pandas version in the virtual env is: {pd.__version__}")
5
6
    num = random.randint(0, 100)
7
8
    if num > 50:
9
        # return the task_id of the downstream task that should be executed
10
        return "downstream_task_a"  
11
    else:
12
        return "downstream_task_b"
13
14
my_isolated_task = BranchPythonVirtualenvOperator(
15
    task_id="my_isolated_task",
16
    python_callable=my_isolated_function,
17
    requirements=["pandas==1.5.1"],
18
)

1
# from airflow.decorators import task
2
# import os
3
4
# note that to be able to use the logical date, pendulum needs to be installed in the epo_pyenv
5
@task.external_python(python=os.environ["ASTRO_PYENV_epo_pyenv"])
6
def my_isolated_task(logical_date): 
7
    print(f"The logical date is: {logical_date}")
8
    # your code to run in the isolated environment 
9
10
my_isolated_task()

1
# from airflow.operators.python import ExternalPythonOperator
2
# import os
3
4
def my_isolated_function(logical_date_from_op_kwargs):
5
    print(f"The logical date is: {logical_date_from_op_kwargs}")
6
    # your code to run in the isolated environment 
7
8
my_isolated_task = ExternalPythonOperator(
9
    task_id="my_isolated_task",
10
    python_callable=my_isolated_function,
11
    # note that to be able to use the logical date, pendulum needs to be installed in the epo_pyenv
12
    python=os.environ["ASTRO_PYENV_epo_pyenv"],
13
    op_kwargs={
14
        "logical_date_from_op_kwargs": "{{ logical_date }}",
15
    },
16
)

1
# from airflow.decorators import task
2
3
@task.virtualenv(
4
    requirements=[
5
        "pandas==1.5.1",
6
        "pendulum==3.0.0",
7
    ],  # pendulum is needed to use the logical date
8
)
9
def my_isolated_task(logical_date):
10
    print(f"The logical date is: {logical_date}")
11
    # your code to run in the isolated environment

1
# from airflow.operators.python import PythonVirtualenvOperator
2
3
def my_isolated_function(logical_date_from_op_kwargs):
4
    print(f"The logical date is: {logical_date_from_op_kwargs}")
5
    # your code to run in the isolated environment 
6
7
my_isolated_task = PythonVirtualenvOperator(
8
    task_id="my_isolated_task",
9
    python_callable=my_isolated_function,
10
    requirements=[
11
        "pandas==1.5.1",
12
        "pendulum==3.0.0",
13
    ],  # pendulum is needed to use the logical date
14
    op_kwargs={
15
        "logical_date_from_op_kwargs": "{{ logical_date }}"
16
    },
17
)

1
# from airflow.operators.python import PythonVirtualenvOperator
2
3
def my_isolated_function(password_from_op_kwargs):
4
    print(f"The password is: {password_from_op_kwargs}")
5
6
my_isolated_task = PythonVirtualenvOperator(
7
    task_id="my_isolated_task",
8
    python_callable=my_isolated_function,
9
    requirements=["pandas==1.5.1"],
10
    python_version="3.10",
11
    op_kwargs={
12
        "password_from_op_kwargs": "{{ var.value.my_secret }}",
13
    },
14
)

1
# from airflow.operators.python import ExternalPythonOperator
2
# import os
3
4
def my_isolated_function(password_from_op_kwargs):
5
    print(f"The password is: {password_from_op_kwargs}")
6
7
my_isolated_task = ExternalPythonOperator(
8
    task_id="my_isolated_task",
9
    python_callable=my_isolated_function,
10
    python=os.environ["ASTRO_PYENV_epo_pyenv"],
11
    op_kwargs={
12
        "password_from_op_kwargs": "{{ var.value.my_secret }}",
13
    },
14
)

1
# from airflow.operators.python import PythonVirtualenvOperator
2
3
def my_isolated_function():
4
    from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
5
    hook = SnowflakeHook(snowflake_conn_id="MY_SNOWFLAKE_CONN_ID")
6
    result = hook.get_first("SELECT * FROM MY_TABLE LIMIT 1")
7
    print(f"The pandas version in the virtual env is: {pd.__version__}")
8
9
    return result
10
11
my_isolated_task = PythonVirtualenvOperator(
12
    task_id="my_isolated_task",
13
    python_callable=my_isolated_function,
14
    requirements=[
15
        "pandas==1.5.3",
16
        "apache-airflow==2.8.1",
17
        "apache-airflow-providers-snowflake==5.3.0",
18
    ],
19
    venv_cache_path="/tmp/venv_cache",  # optional caching of the virtual environment
20
)