Airflow Valohai Plugin

Integration between Airflow and Valohai that allow Airflow tasks to launch executions in Valohai.

Installation

Install this package directly from pypi.

pip install airflow-valohai-plugin

You can read more about the use cases of Airflow plugins in the official docs.

Usage

Get a Valohai API Token.

Create a Valohai Connection in the Airflow UI with:

• Conn Id: valohai_default
• Conn Type: HTTP
• Host: app.valohai.com
• Password: REPLACE_WITH_API_TOKEN

There are two operators that you can import.

from airflow_valohai_plugin.operators.valohai_submit_execution_operator import ValohaiSubmitExecutionOperator
from airflow_valohai_plugin.operators.valohai_download_execution_outputs_operator import ValohaiDownloadExecutionOutputsOperator

You can then create tasks and assign them to your DAGs.

ValohaiSubmitExecutionOperator

train_model = ValohaiSubmitExecutionOperator(
    task_id='train_model',
    project_name='tensorflow-example',
    step='Train model (MNIST)',
    dag=dag,
)

ValohaiDownloadExecutionOutputsOperator

from airflow_valohai_plugin.operators.valohai_download_execution_outputs_operator import ValohaiDownloadExecutionOutputsOperator

download_model = ValohaiDownloadExecutionOutputsOperator(
    task_id='download_model',
    output_task=train_model,
    output_name='model.pb',
    dag=dag
)

And create dependencies between your tasks.

train_model >> download_model

Examples

You can find the complete code of example Airflow DAGs in the examples folder.

Contributing

Check out the contributing page.