Airflow on SherlockML

01_readme.md

Airflow on SherlockML

Apache Airflow is an open source tool for creating task pipelines. It lets you define sets of tasks and dependencies between those tasks, and then takes care of the execution.

Airflow can be a useful add-on to SherlockML: you can schedule periodic jobs that automatically fetch data or retrain a model or publish a report. By defining dependencies between those, you can have a resilient, automated data processing pipeline.

Installing Airflow in SherlockML

Setting up a project

To set up a project, create an environment called airflow-setup-project with the following entry in the scripts tab:

AIRFLOW_HOME=/project/airflow
SOURCE=https://gist.githubusercontent.com/pbugnion/a4a24d6f54fe9f6239aae90569197e9e/raw/

source activate Python3
pip install apache-airflow

mkdir -p $AIRFLOW_HOME

wget --output-document $AIRFLOW_HOME/airflow.cfg $SOURCE/airflow.cfg

echo "export AIRFLOW_HOME=$AIRFLOW_HOME" | sudo tee /etc/profile.d/airflow.sh

airflow initdb

Then apply this environment to any Jupyter server within the project.

Running the Airflow webserver

Create an environment called airflow-webserver with a single shell script:

AIRFLOW_HOME=/project/airflow
SOURCE=https://gist.githubusercontent.com/pbugnion/a4a24d6f54fe9f6239aae90569197e9e/raw

echo "export AIRFLOW_HOME=$AIRFLOW_HOME" | sudo tee /etc/profile.d/airflow.sh

sudo sv stop jupyter

wget --output-document /tmp/airflow-run $SOURCE/run
sudo mkdir -p /etc/service/airflow
sudo mv /tmp/airflow-run /etc/service/airflow/run
sudo chown root:root  /etc/service/airflow/run
sudo chmod a+x /etc/service/airflow/run

sleep 7 # Wait for runit to pick up new service

sudo sv start airflow

sleep 20 # Wait for airflow to start

When this environment is applied to a server in the workspace, it will stop the Jupyter server and replace it with Airflow. You can then click on the server name to open the Airflow dashboard.

Install examples

This gist contains an example with an airflow pipeline that fetches the Google stock price for a given date, saves it in a CSV in the SherlockML workspace (we could also save it in, say, InfluxDB) and publishes a SherlockML report with a graph of the stock price over time, including that date.

To install the demo, run the following commands in a terminal:

SOURCE=https://gist.githubusercontent.com/pbugnion/a4a24d6f54fe9f6239aae90569197e9e/raw

mkdir -p /project/stock-price /project/stock-price/data

wget --output-document /project/stock-price/fetch_stock.py $SOURCE/fetch_stock.py
wget --output-document /project/stock-price/publish_report.py $SOURCE/publish_report.py
wget --output-document /project/stock-price/plot-stock-price.ipynb $SOURCE/plot-stock-price.ipynb

wget --output-document /project/airflow/dags/stock_price.py $SOURCE/stock_price_dag.py

Then create a report by publishing the plot-stock-price notebook in /project/stock-price. Get the report ID for that report (the last UUID in the URL when you are viewing the report). Edit the DAG definition with the correct UUID.

The file /project/airflow/dags/stock_price.py is the airflow DAG definition: it defines what tasks are available in the DAG, and how they depend on each other.

This pipeline is scheduled to run every day at midnight. You can also run the pipeline for a particular day or set of days by running the following command in a new terminal on the instance running the airflow webserver:

airflow backfill stock_price -s 2018-03-01 -e 2018-03-06

This will fetch the stock price for dates between the 1st and the 6th of March 2018.

airflow-example.png

airflow.cfg

# Airflow configuration file
# /project/airflow/airflow.cfg

[core]
# The home folder for airflow, default is ~/airflow
airflow_home = /project/airflow

# The folder where your airflow pipelines live, most likely a
# subfolder in a code repository
# This path must be absolute
dags_folder = /project/airflow/dags

# The folder where airflow should store its log files
# This path must be absolute
base_log_folder = /project/airflow/logs

# Airflow can store logs remotely in AWS S3 or Google Cloud Storage. Users
# must supply a remote location URL (starting with either 's3://...' or
# 'gs://...') and an Airflow connection id that provides access to the storage
# location.
remote_base_log_folder =
remote_log_conn_id =
# Use server-side encryption for logs stored in S3
encrypt_s3_logs = False
# DEPRECATED option for remote log storage, use remote_base_log_folder instead!
s3_log_folder =

# The executor class that airflow should use. Choices include
# SequentialExecutor, LocalExecutor, CeleryExecutor
executor = SequentialExecutor

# The SqlAlchemy connection string to the metadata database.
# SqlAlchemy supports many different database engine, more information
# their website
sql_alchemy_conn = sqlite:////project/airflow/airflow.db

# The SqlAlchemy pool size is the maximum number of database connections
# in the pool.
sql_alchemy_pool_size = 5

# The SqlAlchemy pool recycle is the number of seconds a connection
# can be idle in the pool before it is invalidated. This config does
# not apply to sqlite.
sql_alchemy_pool_recycle = 3600

# The amount of parallelism as a setting to the executor. This defines
# the max number of task instances that should run simultaneously
# on this airflow installation
parallelism = 8

# The number of task instances allowed to run concurrently by the scheduler
dag_concurrency = 16

# Are DAGs paused by default at creation
dags_are_paused_at_creation = True

# When not using pools, tasks are run in the "default pool",
# whose size is guided by this config element
non_pooled_task_slot_count = 128

# The maximum number of active DAG runs per DAG
max_active_runs_per_dag = 16

# Whether to load the examples that ship with Airflow. It's good to
# get started, but you probably want to set this to False in a production
# environment
load_examples = False

# Where your Airflow plugins are stored
plugins_folder = /home/sherlock/workspace/airflow/plugins

# Secret key to save connection passwords in the db
fernet_key = IrY20uqM6R-OGwrRaWFh4rgdh5hD0n2p7a5jgd-nhxw=

# Whether to disable pickling dags
donot_pickle = False

# How long before timing out a python file import while filling the DagBag
dagbag_import_timeout = 30

# The class to use for running task instances in a subprocess
task_runner = BashTaskRunner

# If set, tasks without a `run_as_user` argument will be run with this user
# Can be used to de-elevate a sudo user running Airflow when executing tasks
default_impersonation =

# What security module to use (for example kerberos):
security =

# Turn unit test mode on (overwrites many configuration options with test
# values at runtime)
unit_test_mode = False

[cli]
# In what way should the cli access the API. The LocalClient will use the
# database directly, while the json_client will use the api running on the
# webserver
api_client = airflow.api.client.local_client
endpoint_url = http://localhost:8888

[api]
# How to authenticate users of the API
auth_backend = airflow.api.auth.backend.default

[operators]
# The default owner assigned to each new operator, unless
# provided explicitly or passed via `default_args`
default_owner = Airflow
default_cpus = 1
default_ram = 512
default_disk = 512
default_gpus = 0


[webserver]
# The base url of your website as airflow cannot guess what domain or
# cname you are using. This is used in automated emails that
# airflow sends to point links to the right web server
base_url = http://localhost:8888

# The ip specified when starting the web server
web_server_host = 127.0.0.1

# The port on which to run the web server
web_server_port = 8888

# Paths to the SSL certificate and key for the web server. When both are
# provided SSL will be enabled. This does not change the web server port.
web_server_ssl_cert =
web_server_ssl_key =

# Number of seconds the gunicorn webserver waits before timing out on a worker
web_server_worker_timeout = 120

# Number of workers to refresh at a time. When set to 0, worker refresh is
# disabled. When nonzero, airflow periodically refreshes webserver workers by
# bringing up new ones and killing old ones.
worker_refresh_batch_size = 1

# Number of seconds to wait before refreshing a batch of workers.
worker_refresh_interval = 30

# Secret key used to run your flask app
secret_key = temporary_key

# Number of workers to run the Gunicorn web server
workers = 4

# The worker class gunicorn should use. Choices include
# sync (default), eventlet, gevent
worker_class = sync

# Log files for the gunicorn webserver. '-' means log to stderr.
access_logfile = -
error_logfile = -

# Expose the configuration file in the web server
expose_config = False

# Set to true to turn on authentication:
# http://pythonhosted.org/airflow/security.html#web-authentication
authenticate = False

# Filter the list of dags by owner name (requires authentication to be enabled)
filter_by_owner = False

# Filtering mode. Choices include user (default) and ldapgroup.
# Ldap group filtering requires using the ldap backend
#
# Note that the ldap server needs the "memberOf" overlay to be set up
# in order to user the ldapgroup mode.
owner_mode = user

# Default DAG orientation. Valid values are:
# LR (Left->Right), TB (Top->Bottom), RL (Right->Left), BT (Bottom->Top)
dag_orientation = LR

# Puts the webserver in demonstration mode; blurs the names of Operators for
# privacy.
demo_mode = False

# The amount of time (in secs) webserver will wait for initial handshake
# while fetching logs from other worker machine
log_fetch_timeout_sec = 5

# By default, the webserver shows paused DAGs. Flip this to hide paused
# DAGs by default
hide_paused_dags_by_default = False

[email]
email_backend = airflow.utils.email.send_email_smtp


[smtp]
# If you want airflow to send emails on retries, failure, and you want to use
# the airflow.utils.email.send_email_smtp function, you have to configure an
# smtp server here
smtp_host = localhost
smtp_starttls = True
smtp_ssl = False
# Uncomment and set the user/pass settings if you want to use SMTP AUTH
# smtp_user = airflow
# smtp_password = airflow
smtp_port = 25
smtp_mail_from = [email protected]


[celery]
# This section only applies if you are using the CeleryExecutor in
# [core] section above

# The app name that will be used by celery
celery_app_name = airflow.executors.celery_executor

# The concurrency that will be used when starting workers with the
# "airflow worker" command. This defines the number of task instances that
# a worker will take, so size up your workers based on the resources on
# your worker box and the nature of your tasks
celeryd_concurrency = 16

# When you start an airflow worker, airflow starts a tiny web server
# subprocess to serve the workers local log files to the airflow main
# web server, who then builds pages and sends them to users. This defines
# the port on which the logs are served. It needs to be unused, and open
# visible from the main web server to connect into the workers.
worker_log_server_port = 8793

# The Celery broker URL. Celery supports RabbitMQ, Redis and experimentally
# a sqlalchemy database. Refer to the Celery documentation for more
# information.
broker_url = sqla+mysql://airflow:airflow@localhost:3306/airflow

# Another key Celery setting
celery_result_backend = db+mysql://airflow:airflow@localhost:3306/airflow

# Celery Flower is a sweet UI for Celery. Airflow has a shortcut to start
# it `airflow flower`. This defines the IP that Celery Flower runs on
flower_host = 0.0.0.0

# This defines the port that Celery Flower runs on
flower_port = 5555

# Default queue that tasks get assigned to and that worker listen on.
default_queue = default


[scheduler]
# Task instances listen for external kill signal (when you clear tasks
# from the CLI or the UI), this defines the frequency at which they should
# listen (in seconds).
job_heartbeat_sec = 5

# The scheduler constantly tries to trigger new tasks (look at the
# scheduler section in the docs for more information). This defines
# how often the scheduler should run (in seconds).
scheduler_heartbeat_sec = 5

# after how much time should the scheduler terminate in seconds
# -1 indicates to run continuously (see also num_runs)
run_duration = -1

# after how much time a new DAGs should be picked up from the filesystem
min_file_process_interval = 0

dag_dir_list_interval = 300

# How often should stats be printed to the logs
print_stats_interval = 30

child_process_log_directory = /home/sherlock/workspace/airflow/logs/scheduler

# Local task jobs periodically heartbeat to the DB. If the job has
# not heartbeat in this many seconds, the scheduler will mark the
# associated task instance as failed and will re-schedule the task.
scheduler_zombie_task_threshold = 300

# Turn off scheduler catchup by setting this to False.
# Default behavior is unchanged and
# Command Line Backfills still work, but the scheduler
# will not do scheduler catchup if this is False,
# however it can be set on a per DAG basis in the
# DAG definition (catchup)
catchup_by_default = True

# Statsd (https://github.com/etsy/statsd) integration settings
statsd_on = False
statsd_host = localhost
statsd_port = 8125
statsd_prefix = airflow

# The scheduler can run multiple threads in parallel to schedule dags.
# This defines how many threads will run. However airflow will never
# use more threads than the amount of cpu cores available.
max_threads = 2

authenticate = False


[mesos]
# Mesos master address which MesosExecutor will connect to.
master = localhost:5050

# The framework name which Airflow scheduler will register itself as on mesos
framework_name = Airflow

# Number of cpu cores required for running one task instance using
# 'airflow run <dag_id> <task_id> <execution_date> --local -p <pickle_id>'
# command on a mesos slave
task_cpu = 1

# Memory in MB required for running one task instance using
# 'airflow run <dag_id> <task_id> <execution_date> --local -p <pickle_id>'
# command on a mesos slave
task_memory = 256

# Enable framework checkpointing for mesos
# See http://mesos.apache.org/documentation/latest/slave-recovery/
checkpoint = False

# Failover timeout in milliseconds.
# When checkpointing is enabled and this option is set, Mesos waits
# until the configured timeout for
# the MesosExecutor framework to re-register after a failover. Mesos
# shuts down running tasks if the
# MesosExecutor framework fails to re-register within this timeframe.
# failover_timeout = 604800

# Enable framework authentication for mesos
# See http://mesos.apache.org/documentation/latest/configuration/
authenticate = False

# Mesos credentials, if authentication is enabled
# default_principal = admin
# default_secret = admin


[kerberos]
ccache = /tmp/airflow_krb5_ccache
# gets augmented with fqdn
principal = airflow
reinit_frequency = 3600
kinit_path = kinit
keytab = airflow.keytab


[github_enterprise]
api_rev = v3


[admin]
# UI to hide sensitive variable fields when set to True
hide_sensitive_variable_fields = True

fetch_stock.py

import json
import csv
import logging
from pathlib import Path
import argparse
from collections import namedtuple

import requests

logging.basicConfig(level=logging.INFO)

QUANDL_URL = 'https://www.quandl.com/api/v3/datasets/WIKI/{}.json'
DATASTORE = Path('/project/stock-price/data')

Configuration = namedtuple(
    'Configuration',
    ['stock', 'start_date', 'end_date']
)

def fetch_stock_info(stock, start_date, end_date):
    response = requests.get(
        QUANDL_URL.format(stock),
        params={'start_date': start_date, 'end_date': end_date}
    )
    response.raise_for_status()
    json = response.json()
    data = json['dataset']['data']
    stock_price = {}
    for entry in data:
        date = entry[0]
        closing = entry[4]
        stock_price[date] = float(closing)
    return stock_price
    

def read_current_data(store):
    stock_prices = {}
    try:
        with open(store) as f:
            reader = csv.reader(f)
            next(reader)  # skip header
            for row in reader:
                date, price = row
                stock_prices[date] = price
    except FileNotFoundError:
        return {}
    return stock_prices
    
    
def write_data(store, stock_prices):
    sorted_data = sorted(stock_prices.items(), key=lambda it: it[0])
    with open(store, 'w') as f:
        writer = csv.writer(f)
        writer.writerow(['date', 'price'])
        writer.writerows(sorted_data)
    
    
def store_update(store, new_prices):
    stock_prices = read_current_data(store)
    stock_prices.update(new_prices)
    write_data(store, stock_prices)
    
    
def fetch_and_update_stocks(stock, start_date, end_date):
    logging.info(
        f'Fetching stocks for symbol {stock} '
        f'between {start_date} and {end_date}'
    )
    prices = fetch_stock_info(stock, start_date, end_date)
    logging.info(f'Fetched {len(prices)} stock prices.')
    store = DATASTORE / f'{stock}.csv'
    store_update(store, prices)
    logging.info(f'Written information back to data store {store}.')
    
    
def parse_command_line():
    parser = argparse.ArgumentParser(
        description='Fetch and store stock information')
    parser.add_argument('symbol')
    parser.add_argument(
        'start_date',
        help='start date, in format 2018-01-22'
    )
    parser.add_argument(
        'end_date',
        help='end date, in format 2018-01-29'
    )
    args = parser.parse_args()
    return Configuration(
        args.symbol.upper(),
        args.start_date,
        args.end_date
    )
    
    
if __name__ == '__main__':
    configuration = parse_command_line()
    logging.info(f'Using configuration {configuration}.')
    fetch_and_update_stocks(
        configuration.stock, 
        configuration.start_date,
        configuration.end_date
    )
    

plot-stock-price.ipynb

publish_report.py

import os
import argparse

import requests

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor

import sml.auth
import sml.config


HUDSON_URL = sml.config.url_for_service('hudson')
TAVERN_URL = sml.config.url_for_service('tavern')


class InputError(Exception):
    pass


def run(report_id, notebook, run_notebook=True):
    """Publish a report in the current project.

    Parameters
    ----------
    report_id: str
        SherlockML report ID.
    notebook: str
        Full path to notebook on SherlockML.
    run_notebook: bool, optional
        Run notebook before publishing (default: True).
    """

    # Run the notebook
    if run_notebook:
        with open(notebook) as f:
            nb = nbformat.read(f, as_version=4)

        ep = ExecutePreprocessor(kernel_name='Python3')
        ep.preprocess(nb, {'metadata': {'path': os.path.dirname(notebook)}})

        with open(notebook, 'wt') as f:
            nbformat.write(nb, f)

    # Make sure notebook path starts with /project
    split_notebook_path = notebook.split('/')
    if split_notebook_path[0] == '' and split_notebook_path[1] == 'project':
        pass
    else:
        raise InputError('notebook path must start with /project')

    # Get rid of the project bit
    reconstructed_notebook_path = []
    for i, item in enumerate(split_notebook_path):
        if i < 2:
            continue
        reconstructed_notebook_path.append(item)
    reconstructed_notebook_path = '/'.join(reconstructed_notebook_path)

    # Authenticate and publish
    auth_headers = sml.auth.auth_headers()
    user_id_response = requests.get(
        HUDSON_URL + '/authenticate',
        headers=auth_headers
    )
    user_id_response.raise_for_status()
    user_id = user_id_response.json()['account']['userId']
    body = {
        'notebook_path': reconstructed_notebook_path,
        'author_id': user_id,
        'draft': False
    }
    create_version_response = requests.post(
        TAVERN_URL + '/report/{}/version'.format(report_id),
        json=body,
        headers=auth_headers)
    create_version_response.raise_for_status()


def parse_command_line():
    parser = argparse.ArgumentParser(
        description='Run and publish a notebook')
    parser.add_argument('report_id', help='ID of the report to update')
    parser.add_argument(
        'notebook_path',
        help='Path of the notebook to publish'
    )
    args = parser.parse_args()
    return (args.report_id, args.notebook_path)


if __name__ == '__main__':
    report_id, notebook_path = parse_command_line()
    run(report_id, notebook_path)

report-example.png

run

#!/bin/sh                                                

# Runit unit file for Airflow webserver

set -e                                                   
                                                         
cd "$PROJECT_HOME"                                       
                                                         
exec /sbin/setuser sherlock env \
  AIRFLOW_HOME=/project/airflow \
  PATH=/opt/anaconda/envs/Python3/bin:opt/anaconda/bin:"$PATH" \
  airflow webserver

stock_price_dag.py

from airflow import DAG
from airflow.operators.bash_operator import BashOperator
from datetime import datetime, timedelta

REPORT_ID = 'b0b5e1c5-d7c9-4ae5-8b02-ce0f4a7ddaee'

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2017, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
    'schedule_interval': '@daily'
}

dag = DAG(
    'stock_price',
    default_args=default_args
)


fetch_command = (
    '/opt/anaconda/envs/Python3/bin/python '
    '/project/stock-price/fetch_stock.py GOOG {{ ds }} {{ ds }}'
)

fetch_operator = BashOperator(
    task_id='fetch_price',
    bash_command=fetch_command,
    dag=dag
)

publish_operator = BashOperator(
    task_id='publish',
    bash_command=(
        '/opt/anaconda/envs/Python3/bin/python '
        '/project/stock-price/publish_report.py '
        f'{REPORT_ID} '
        '/project/stock-price/plot-stock-price.ipynb'
    ),
    dag=dag
)


publish_operator.set_upstream(fetch_operator)