Skip to content

🦈 Airflow plugin to scale machine learning tasks with Valohai and get automatic version control

License

Notifications You must be signed in to change notification settings

skillupco/airflow-valohai-plugin

Repository files navigation

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.