if [ "$AIRFLOW__CORE__EXECUTOR" = "CeleryExecutor" ]; then
  if [ "$USE_SQS_FOR_BROKER" != "" ]; then
    echo "use sqs for broker"
    AIRFLOW__CELERY__BROKER_URL="sqs://"
  else
    AIRFLOW__CELERY__BROKER_URL="redis://$REDIS_PREFIX$REDIS_HOST:$REDIS_PORT/1"
    wait_for_port "Redis" "$REDIS_HOST" "$REDIS_PORT"
  fi
fi

a >> b
## b << a　でもOK

a >> b >> c

a >> [b >> c] >> d

a >> [b1 >> b2, c] >> d

def wrap(op):
    if not isinstance(op, list):
        op = [op]
    return OperatorWrapper(op, op)
​
​
class OperatorWrapper:
    def __init__(self, lefts, rights):
        self.lefts = [b for a in lefts for b in (a.lefts if isinstance(a, OperatorWrapper) else [a])]
        self.rights = [b for a in rights for b in (a.rights if isinstance(a, OperatorWrapper) else [a])]
​
    def __rshift__(self, other):
        if not isinstance(other, OperatorWrapper):
            other = wrap(other)
        for lop in self.rights:
            for rop in other.lefts:
                lop >> rop
        return OperatorWrapper(self.lefts, other.rights)
​
    def __lshift__(self, other):
        if not isinstance(other, OperatorWrapper):
            other = wrap(other)
        for lop in self.rights:
            for rop in other.lefts:
                lop << rop
        return OperatorWrapper(self.lefts, other.rights)
​
    def __rrshift__(self, other):
        return (self >> other)._reversed()
​
    def __rlshift__(self, other):
        return (self << other)._reversed()
​
    def _reversed(self):
        return OperatorWrapper(self.rights, self.lefts)

wrap(a) >> [wrap(b1) >> b2, c] >> d

def main(params, bash_cmd, angles):
    today = datetime.now()
    month_ago = today-relativedelta(months=1)
    today_ops = [
        BashOperator(
            task_id=f'sync_data_today_angle_{angle}',
            bash_command=bash_cmd,
            xcom_push=True,
            params={
                    **params,
                    'year': today.year,
                    'month': str(today.month).zfill(2),
                    'angle': angle,
            },
        )
        for angle in angles
    ]
    month_ago_ops = [
        BashOperator(
            task_id=f'sync_data_month_ago_angle_{angle}',
            bash_command=bash_cmd,
            xcom_push=True,
            params={
                    **params,
                    'year': month_ago.year,
                    'month': str(month_ago.month).zfill(2),
                    'angle': angle,
            },
        )
        for angle in angles
    ]
    report_operator = PythonOperator(
        task_id='report',
        provide_context=True,
        python_callable=reporter('return_value'),
    )
    return wrap(today_ops + month_ago_ops) >> report_operator

with DAG(
    'sync_data_main',
    description='本番dag',
    default_args=default_args,
    max_active_runs=1,
    schedule_interval='@hourly',
    catchup=False
) as sync_data_main:
    params = {
        'src_bucket': os.environ.get('SRC_BUCKET'),
        'dst_bucket': os.environ.get('DST_BUCKET'),
        'queue_name': os.environ.get('QUEUE_NAME'),
        'env': 'prod',
    }
    sync_data_main >> main(params, bash_cmd, angles)

with DAG(
    'sync_data_dev',
    description='開発用。コピー先バケットとsqsキューが異なる',
    default_args=default_args,
    max_active_runs=1,
    schedule_interval=None,
    catchup=False
) as sync_data_dev:
    params = {
        'src_bucket': os.environ.get('SRC_BUCKET_DEV'),
        'dst_bucket': os.environ.get('DST_BUCKET_DEV'),
        'queue_name': os.environ.get('QUEUE_NAME_DEV'),
        'env': 'dev',
    }
    sync_data_dev >> main(params, bash_cmd, angles)

## serving.py
from sagemaker_containers.beta.framework import worker
from model import Model # 同じディレクトリにあるmodel.py


## モデルの読み込み。
def model_fn(model_dir):
    model_path = os.path.join(model_dir, 'params.npz')
    config = make_model_config(model_dir)
    return Model(config, model_path)

## 入力データを受け取る。
## endpointの場合httpリクエストペイロードが、batchtransformの場合s3に置かれたファイルの中身が引数として渡される。
## その違いを意識する必要はない。
def input_fn(input_bytes, content_type):
    if content_type == 'application/json':
        # データパスを書いたjsonを受け取る場合
        req_body = json.loads(input_bytes)
        bucket = req_body['bucket']

        data = get_data_from_s3(bucket, req_body['data_path'])
        return data
    else:
        # 生データを受け取る場合
        data = load_as_numpy_array(input_bytes)
        return data

## 推論の実行。model_fnとinput_fnの結果が渡される。
def predict_fn(data, model):
    prediction = model.predict(data)
    return prediction

## レスポンスの生成
## 結果はendpointの場合レスポンスbodyに、batchtransformの場合s3にファイルとして書き出される。
## ここでも違いを意識する必要はない
def output_fn(prediction, accept):
    with BytesIO() as f:
        np.save(f, prediction)
        return worker.Response(response=f.getvalue(), mimetype=accept) # 単にバイト列をreturnするだけでもいけた気がする

import boto3
import sagemaker
from sagemaker import Session
from sagemaker.local import LocalSession
from sagemaker.chainer.model import ChainerModel
from sagemaker.utils import sagemaker_timestamp

from utils import DataTransfer


def create_model(sagemaker_session, config, model_name, region_name):
    parameters = {
        'sagemaker_session': sagemaker_session,
        'model_data': config['model']['path'],
        'role': config['role'], # 実行用 IAM role
        'env': config['env'],
        'entry_point': 'serving.py'
        'source_dir': 'models/hoge/', # serving.pyがあるディレクトリ
        'dependencies': 'utils/', # source_dir以外の依存コードがあるディレクトリ
        'vpc_config': get_vpc_config(region_name), # サブネット設定など
        'framework_version': '5.0.0',
        'py_version': 'py3',
        'model_server_workers': 1,
        'name': model_name,
    }
    return ChainerModel(**parameters)


def predict(model_name, data_infos, instance_count, bucket_name, mode, region_name):
    data = [json.dumps(datainfo, indent=2) for datainfo in data_infos]

    boto3_session = boto3.Session(region_name=region_name)
    s3 = boto3_session.client('s3')

    data_transfer = DataTransfer(s3, bucket_name, 'tmp/hoge')
    data_transfer.send(data)

    sagemaker_config = get_sagemaker_config()
    if mode == 'aws':
        sagemaker_session = Session(boto_session=boto3_session)
    elif mode == 'local':
        sagemaker_session = LocalSession()

    m = create_model(
        sagemaker_session,
        sagemaker_config[mode],
        model_name,
        region_name,
    )
    transformer = m.transformer(
        instance_count,
        sagemaker_config[mode]['instance_type'],
        accept='image/jpeg',
        output_path=f's3://{bucket_name}/{data_transfer.output_dir}',
        strategy='SingleRecord',
        assemble_with=None,
        output_kms_key=None,
        env=None,
        max_concurrent_transforms=None,
        max_payload=None,
        tags=None,
        volume_kms_key=None,
    )
    transformer.transform(
        f's3://{bucket_name}/{data_transfer.input_dir}',
        job_name=None,
        data_type='S3Prefix',  # or "ManifestFile"
        content_type='application/json',
        compression_type=None,  # or 'Gzip'
        split_type=None,  # 'None', 'Line', 'RecordIO', 'TFRecord'
        input_filter=None,  # "$[1:]", "$.features"
        output_filter=None,  # "$[1:]", "$.prediction"
        join_source=None,  # Input, None
        wait=True
    )

    return [
        np.load(BytesIO(it), allow_pickle=True)
        for it in data_transfer.receive()
    ]

def deploy_endpoint(sagemaker_session, config, endpoint_name, region):
    model = create_model(sagemaker_session, config, f'hoge-{endpoint_name}', region)

    # deploy model to endpoint
    model.deploy(
        instance_type=config['instance_type'],
        initial_instance_count=1,
        endpoint_name=endpoint_name
    )

## 再掲
with DAG(
    'sync_data_main',
    description='本番dag',
    default_args=default_args,
    max_active_runs=1, # 平行実行されないように
    schedule_interval='@hourly', # 1時間おきに実行
    catchup=False # 実行していない過去のスケジュールをskip
) as sync_data_main:

from airflow.sensors.base_sensor_operator import BaseSensorOperator


class SQSBufferingSensor(BaseSensorOperator):
    """
    標準のSQSSensorはメッセージを最大10しか取得できないので、無制限に取得できるようにした
    """

    def __init__(self,
                 sqs_queue_url,
                 max_messages,
                 max_wait_seconds,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        self.sqs_queue_url = sqs_queue_url
        self.start_time = None
        self.message_buffer = []
        self.max_messages = max_messages
        self.max_wait_seconds = max_wait_seconds
        self.session = boto3.Session()

    def poke(self, context):
        if self.start_time is None:
            self.start_time = time.time()

        sqs = self.session.client('sqs')

        response = sqs.receive_message(
            QueueUrl=self.sqs_queue_url,
            MaxNumberOfMessages=min(10, self.max_messages - len(self.message_buffer)),
            WaitTimeSeconds=5,
        )

        if 'Messages' in response and len(response['Messages']) > 0:
            messages = response['Messages']
            entries = [
                {
                    'Id': message['MessageId'],
                    'ReceiptHandle': message['ReceiptHandle']
                }
                for message in messages
            ]
            self.message_buffer.extend(messages)

        elapsed = time.time()-self.start_time
        if len(self.message_buffer) >= self.max_messages:
            # max_messages以上のメッセージを受信したら終了
            context['ti'].xcom_push(key='messages', value={'Messages': self.message_buffer})
            return True

        if elapsed > self.max_wait_seconds:
            # max_wait_seconds経過したら数が足りなくても終了
            context['ti'].xcom_push(key='messages', value={'Messages': self.message_buffer})
            return True

        return False

from airflow.plugins_manager import AirflowPlugin
from sensors import SQSBufferingSensor


class HogePlugin(AirflowPlugin):
    name = 'hoge_plugin'
    operators = []
    sensors = [SQSBufferingSensor]
    hooks = []
    executors = []
    macros = []
    admin_views = []
    flask_blueprints = []
    menu_links = []
    appbuilder_views = []
    appbuilder_menu_items = []
    global_operator_extra_links = []

from airflow.sensors.hoge_plugin import SQSBufferingSensor

## あるモデルを呼び出すタスク例
## 他のモデルも大体同じ
def main(params, **context):

    # 引数は params で渡す。
    # params は Airflow のタスク詳細画面から確認できるので管理しやすい。
    BUCKET = params.get('bucket')
    INSTANCE_COUNT = params.get('instance_count')
    SENDER_ID = params.get('xcom_sender_taskid') # ミニバッチを送出したタスクのid
    REGION = params.get('region')

    # xcomからデータを取得
    data = context['ti'].xcom_pull(task_ids=SENDER_ID)

    # overwriteまたはファイルがないときだけ処理
    request_payload = [
        {
            'bucket': BUCKET,
            'data1_path': get_data1_path(d),
            'data2_path': get_data2_path(d),
            'output_path': get_output_path(d),
        }
        for d in data
        if d.get('overwrite') or not s3_file_exists(BUCKET, get_output_path(d))
    ]

    if len(request_payload) == 0:  # 空だったら何もしない
        # タスクが大量にあるとログが一覧しにくくてデバッグが困難なので、ログを集約するタスクを用意している(後述)
        report_message(context['ti'], 'result file is already exists. do nothing.')
        return []

    return predict(
        'airflow-hoge',
        request_payload,
        INSTANCE_COUNT,
        BUCKET,
        'aws',
        REGION,
    )

## 主要なモデルの依存定義
def extract_feat(xcom_sender_taskid, bucket, region):
    return (
        wrap(PythonOperator(
            task_id='model1',
            provide_context=True,
            python_callable=model1.main,
            params={
                'xcom_sender_taskid': xcom_sender_taskid,
                'bucket': bucket,
                'instance_count': get_instance_count('model1'),
                'region': region,
            },
        ))
        >> PythonOperator(
            task_id='model2',
            provide_context=True,
            python_callable=model2.main,
            params={
                'xcom_sender_taskid': xcom_sender_taskid,
                'bucket': bucket,
                'instance_count': get_instance_count('model2'),
                'region': region,
            },
        )
        >> [
            extract_other_feat(
                xcom_sender_taskid=xcom_sender_taskid,
                bucket=bucket,
                region=region,
            ),
            PythonOperator(
                task_id='model3',
                provide_context=True,
                python_callable=model3.main,
                params={
                    'xcom_sender_taskid': xcom_sender_taskid,
                    'bucket': bucket,
                    'instance_count': get_instance_count('model3'),
                    'region': region,
                }
            ),
        ]
    )

## 推論処理全体
def main_ops(xcom_sender_taskid, bucket, region):
    return [
        extract_feat(
            xcom_sender_taskid=xcom_sender_taskid,
            bucket=bucket,
            region=region,
        ),
        PythonOperator(
            task_id='model4',
            provide_context=True,
            python_callable=model4.main,
            params={},
        ),
        PythonOperator(
            task_id='model5',
            provide_context=True,
            python_callable=model5.main,
            params={},
        ),
    ]

with DAG(
    'main',
    default_args=default_args,
    catchup=False,
    schedule_interval=timedelta(seconds=60 * 10), # 10分おきに実行
    max_active_runs=1
) as main:
    (
        main
        # SQSのポーリング
        >> wrap(SQSBufferingSensor(
            sqs_queue_url=INPUT_QUEUE,
            max_messages=100,
            max_wait_seconds=2 * 60,
            task_id='wait_sqs',
            poke_interval=10,
            region=SAGEMAKER_REGION,
        ))
        # メッセージの整形
        >> PythonOperator(
            task_id='parse_message',
            provide_context=True,
            python_callable=parse_message,
        )
        # メッセージが空だった場合以降の処理をskip
        >> ShortCircuitOperator(
            task_id='skip_if_messages_empty',
            provide_context=True,
            python_callable=skip_if_messages_empty,
        )
        # 推論処理本体
        >> main_ops(
            xcom_sender_taskid='parse_message',
            bucket=BUCKET,
            region=SAGEMAKER_REGION,
        )
        # 全タスクが正常終了していたら実行されるSQSメッセージの削除処理
        >> PythonOperator(
            task_id='delete_message',
            provide_context=True,
            python_callable=delete_message,
            params={
                'xcom_sender_taskid': 'wait_sqs',
                'sqs_queue_url': INPUT_QUEUE,
                'sqs_queue_region': SAGEMAKER_REGION,
            }
        )
        # ログの集約
        >> PythonOperator(
            task_id='report',
            provide_context=True,
            python_callable=reporter(),
        )
    )

    # overwriteまたはファイルがないときだけ処理
    request_payload = [
        {
            'bucket': BUCKET,
            'data1_path': get_data1_path(d),
            'data2_path': get_data2_path(d),
            'output_path': get_output_path(d),
        }
        for d in data
        if d.get('overwrite') or not s3_file_exists(BUCKET, get_output_path(d))
    ]

## 自身の上流タスクを再帰的にすべて取得
def get_all_upstream_recursive(task):
    res = []
    processed_task_ids = set()
    currents = [task]
    while len(currents) != 0:
        news = []
        for c in currents:
            for u in c.upstream_list:
                if u.task_id in processed_task_ids:
                    continue
                processed_task_ids.add(u.task_id)
                news.append(u)
                res.append(u)
        currents = news
    return res


def reporter(key: str = 'report'):
    def _op(**context):
        task_ids = [a.task_id for a in get_all_upstream_recursive(context['task'])]
        reports = context['ti'].xcom_pull(task_ids=task_ids, key=key)
        for id, report in zip(task_ids, reports):
            if report is None:
                print(f'no report from {id}')
                continue
            print(f"--------------report from {id}------------------------")
            print(json.dumps(report, indent=2))
    return _op


def report_message(ti, message):
    ti.xcom_push(key='report', value=message)