Apache Airflow虚拟envs变得简单

让在Dockerfile中运行隔离的Python虚拟环境（venv）中的任务变得简单。由Astronomer用❤️维护。

假设你想要能够运行针对Snowflake的Snowpark的Airflow任务 -- 这需要Python 3.8。

在Airflow 2.4中添加了ExternalPythonOperator之后，这是可能的，但管理构建过程以获得干净的、快速的Docker构建可能需要大量的管道。

此存储库提供了解决此问题的优秀打包解决方案，并且与Docker镜像缓存配合良好。

概要

创建一个requirements.txt文件

例如，snowpark-requirements.txt

snowflake-snowpark-python[pandas]

# To get credentials out of a connection we need these in the venv too sadly
apache-airflow
psycopg2-binary
apache-airflow-providers-snowflake

使用我们的自定义Docker构建前端

# syntax=quay.io/astronomer/airflow-extensions:v1

FROM quay.io/astronomer/astro-runtime:7.2.0-base

PYENV 3.8 snowpark snowpark-requirements.txt

注意：那个第一个# syntax=注释很重要，不要省略它！

有关新PYENV指令的更多信息

在DAG中使用它

from __future__ import annotations

import sys

from airflow import DAG
from airflow.decorators import task
from airflow.utils.timezone import datetime

with DAG(
    dag_id="astro_snowpark",
    schedule=None,
    start_date=datetime(2022, 1, 1),
    catchup=False,
    tags=["example"],
) as dag:

    @task
    def print_python():
        print(f"My python version is {sys.version}")

    @task.venv("snowpark")
    def snowpark_task():
        from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
        from snowflake.snowpark import Session

        print(f"My python version is {sys.version}")

        hook = SnowflakeHook("snowflake_default")
        conn_params = hook._get_conn_params()
        session = Session.builder.configs(conn_params).create()
        tables = session.sql("show tables").collect()
        print(tables)

        df_table = session.table("sample_product_data")
        print(df_table.show())
        return df_table.to_pandas()

    @task
    def analyze(df):
        print(f"My python version is {sys.version}")
        print(df.head(2))

    print_python() >> analyze(snowpark_task())

要求

这需要Apache Airflow 2.4+以使ExternalPythonOperator工作。

构建Docker镜像的要求

这需要Docker的buildkit后端。

对于Docker Desktop用户，默认启用；Linux用户需要启用它

在运行docker build命令时设置BuildKit环境变量，请运行

DOCKER_BUILDKIT=1 docker build .

要默认启用Docker BuildKit，请将daemon配置文件位于/etc/docker/daemon.json中的feature设置为true并重启守护进程。如果daemon.json文件不存在，则创建一个名为daemon.json的新文件，然后添加以下内容到该文件。

{
  "features": {
    "buildkit" : true
  }
}

然后重启Docker守护进程。

语法扩展还期望在Docker上下文目录中找到packages.txt和requirements.txt文件（默认情况下这些文件可以为空）。

参考

PYENV Docker指令

PYENV命令向docker镜像添加一个运行在指定Python版本的Python虚拟环境，并可选择从requirements.txt文件安装包。

其语法如下

PYENV <python-version> <venv-name> [<reqs-file>]

需求文件是可选的，因此可以像这样安装一个裸Python环境

PYENV 3.10 venv1

@task.venv装饰器

待办事项！编写装饰器，然后填写文档！

在本仓库中

buildkit/

这包含自定义Docker BuildKit前端（有关详细信息，请参阅此博客），添加了一个新的自定义命令PYENV，可以在Dockerfile中使用它来安装新的Python版本和带有自定义依赖项的虚拟环境。

provider/

这包含一个Apache Airflow提供程序，提供@task.venv装饰器。

详细信息

也就是如何手动完成这一切？

# syntax行告诉buildkit使用我们的Build前端处理Dockerfile中的指令。

上面的示例Dockerfile被转换成以下大致指令

USER root
COPY --link --from=python:3.8-slim /usr/local/bin/*3.8* /usr/local/bin/
COPY --link --from=python:3.8-slim /usr/local/include/python3.8* /usr/local/include/python3.8
COPY --link --from=python:3.8-slim /usr/local/lib/pkgconfig/*3.8* /usr/local/lib/pkgconfig/
COPY --link --from=python:3.8-slim /usr/local/lib/*3.8*.so* /usr/local/lib/
COPY --link --from=python:3.8-slim /usr/local/lib/python3.8 /usr/local/lib/python3.8
RUN /sbin/ldconfig /usr/local/lib
RUN ln -s /usr/local/include/python3.8 /usr/local/include/python3.8m

USER astro
RUN mkdir -p /home/astro/.venv/snowpark
COPY reqs/venv1.txt /home/astro/.venv/snowpark/requirements.txt
RUN /usr/local/bin/python3.8 -m venv --system-site-packages /home/astro/.venv/snowpark
ENV ASTRO_PYENV_snowpark /home/astro/.venv/snowpark/bin/python
RUN --mount=type=cache,target=/home/astro/.cache/pip /home/astro/.venv/snowpark/bin/pip --cache-dir=/home/astro/.cache/pip install -r /home/astro/.venv/snowpark/requirements.txt

Airflow操作符的这一部分的最终任务是查找创建的venv中python的路径，使用ASTRO_PYENV_*环境变量

@task.external_python(python=os.environ["ASTRO_PYENV_snowpark"])
def snowpark_task():
    ...