fren
← writing
24 January 2023

SageMaker and EMR from Airflow: an orchestrator that computes nothing

A managed DAG has neither the memory nor the permissions to carry a training run. It triggers, it waits, it wakes somebody, and that last part is always written last, usually after the incident.

The temptation on a managed Airflow is to run the computation inside the DAG: a `PythonOperator` that loads a dataset and trains a model. It works on a sample and never afterwards. A managed service’s workers are sized to schedule, not to compute; they share memory with the scheduler, and a training run that swells takes the scheduling of every other chain down with it.

The rule I have applied since fits in a sentence: the orchestrator computes nothing. It describes a job, hands it to a service built to carry it, and watches. Everything else (machine sizes, permissions, the cluster’s lifetime) becomes configuration rather than code.

01Describe, trigger, wait

Two shapes recur. Training goes to the model service, with its own image and its own instance type. Heavy preparation goes to an ephemeral cluster that is created, loaded with steps, and shut down: the last operator being the one people forget, and the most expensive one to forget.

python
from airflow.providers.amazon.aws.operators.emr import (
    EmrAddStepsOperator, EmrCreateJobFlowOperator, EmrTerminateJobFlowOperator,
)
from airflow.providers.amazon.aws.operators.sagemaker import SageMakerTrainingOperator
from airflow.providers.amazon.aws.sensors.emr import EmrStepSensor

cluster = EmrCreateJobFlowOperator(task_id="cluster", job_flow_overrides=EMR_CONFIG)

prepare = EmrAddStepsOperator(
    task_id="prepare",
    job_flow_id=cluster.output,
    steps=SPARK_STEPS,
)

# The sensor is what makes the dependency true: without it the next task starts
# as soon as the step is *submitted*, not when it has finished.
wait = EmrStepSensor(
    task_id="wait",
    job_flow_id=cluster.output,
    step_id=prepare.output[0],
    mode="reschedule",  # gives the slot back instead of holding it
    poke_interval=60,
)

training = SageMakerTrainingOperator(
    task_id="training",
    config=TRAINING_CONFIG,
    wait_for_completion=True,
)

shutdown = EmrTerminateJobFlowOperator(
    task_id="shutdown",
    job_flow_id=cluster.output,
    trigger_rule="all_done",  # shut down even if the preparation failed
)

cluster >> prepare >> wait >> training
wait >> shutdown

Two details are worth the paragraph they take. A sensor in `reschedule` mode gives its slot back between checks instead of holding it for two hours; on a managed service where slots are counted and billed, the difference shows on the invoice as much as on throughput. And the `all_done` trigger rule on the shutdown is the only thing guaranteeing that a cluster costing several euros an hour does not outlive the step it served.

02What a managed service takes away

On a managed Airflow you choose neither the version of your dependencies nor the moment the environment restarts. Packages are declared in a file dropped on object storage, and any change triggers an environment restart: tens of minutes during which nothing is scheduled. The practical consequence is that you do not add a dependency lightly, and you prefer a provider operator to a library that would have to be installed.

03Waking somebody, and not for nothing

A failure email is not an alert: nobody reads it at three in the morning. The chain feeding a model in production deserves an on-call page, and the others do not. The sorting is what matters, not the tool: a team woken for a chain with no consequence stops answering the real ones.

python
from airflow.providers.pagerduty.hooks.pagerduty_events import PagerdutyEventsHook

def page_on_call(context) -> None:
    instance = context["task_instance"]
    PagerdutyEventsHook(pagerduty_events_conn_id="pagerduty").create_event(
        summary=f"{instance.dag_id}.{instance.task_id} failed",
        severity="critical",
        source="airflow",
        # The dedup key avoids opening one incident per attempt: the three tries
        # of a single run wake exactly one person.
        dedup_key=f"{instance.dag_id}:{instance.task_id}:{context['run_id']}",
        custom_details={"log": instance.log_url},
    )

# Set on the task that matters, not on the whole DAG: the sorting is what keeps
# an on-call rota credible.
training.on_failure_callback = page_on_call
04What I leave out

I do not compare orchestrators. The subject is interesting and it cannot be settled from one engagement: what counted here was the constraint of the managed service, and it is the same whatever opinion you hold about the tool. Nor do I say anything about resuming an interrupted training run: we relaunched from the start, which was bearable at that scale and would not have been beyond it.