Apache Airflow: How to run a Docker Operator with environment variables from another task?

user_3966019 0
Viewed 4171

I want to run a Docker Operator in Airflow with an environment variable download_path that gets set in a previous task. How can I achieve that? Via Xcom? Minimal example:

# define python function
def make_folder(folder_path:str, date:str):
    download_path= folder_path + date 
    return(download_path)
    os.mkdir(download_path)

# Python operator
task_1 = PythonOperator(
    task_id="make_folder",
    provide_context=False,
    python_callable=make_folder,
    op_kwargs={'folder_path': '/my_path_', 'date': 'str(datetime.date(datetime.today()))'},
    xcom_push=True,
    xcom_all=True
)

# docker operator with download_path as a necessary env variable
task_2 = DockerOperator(
    task_id='docker',
    image='file_processor:latest',
    api_version='auto',
    auto_remove=False,
    command='',
    environment={
        'DPATH': **download_path**
    },
    docker_url="unix://var/run/docker.sock",
    network_mode="bridge",
    xcom_push=True,
    xcom_all=True
)

task_1 >> task_2
airflow docker
2 Answers

Yes, you need to use Xcom. Try this:

task_2 = DockerOperator(
                task_id='docker',
                image='file_processor:latest',
                api_version='auto',
                auto_remove=False,
                command='',
                environment={
                        'DPATH': '{{ti.xcom_pull(task_ids='make_folder') }}'
                },
                docker_url="unix://var/run/docker.sock",
                network_mode="bridge",
                xcom_push=True,
                xcom_all=True
)

If the above doesn't work pass the task id of previous task (make_folder) to the operator and try to pull from there:

context['ti'].xcom_pull(prev_task_id)
user_94981250

You would use xcom indeed. The task_instance object provides the two handful methods for this purpose : xcom_push and xcom_pull.

To push the value to xcom, you need to provide the context to your "python collable" function. This will make the task_instance object accessible to the function.

The environment field of DockerOperator is templated. Which means that it supports macros. You can use the {{ task_instance }} macro.

def make_folder(folder_path:str, date:str, **context):
    download_path= folder_path + date 
    os.mkdir(download_path)
    task_instance = context['task_instance']
    task_instance.xcom_push(key="download_path", value=download_path)


task_1 = PythonOperator(
    task_id="make_folder",
    provide_context=True,
    python_callable=make_folder,
    op_kwargs={'folder_path': '/my_path_', 'date':'str(datetime.date(datetime.today()))'}
)


task_2 = DockerOperator(
                task_id='docker',
                image='file_processor:latest',
                api_version='auto',
                auto_remove=False,
                command='',
                environment={
                  'DPATH': "{{ task_instance.xcom_pull(task_ids='make_folder', key='download_path') }}"
                },
                docker_url="unix://var/run/docker.sock",
                network_mode="bridge"
)

task_1 >> task_2
user_55546660
Related