Launch a Pod in an EKS cluster on AWS

1
# import DAG object and utility packages
2
from airflow import DAG
3
from pendulum import datetime
4
from airflow.configuration import conf
5
6
# import the KubernetesPodOperator
7
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import (
8
    KubernetesPodOperator,
9
)
10
11
# import EKS related packages from the Amazon Provider
12
from airflow.providers.amazon.aws.hooks.eks import EksHook, NodegroupStates
13
from airflow.providers.amazon.aws.operators.eks import (
14
    EksCreateNodegroupOperator,
15
    EksDeleteNodegroupOperator,
16
)
17
from airflow.providers.amazon.aws.sensors.eks import EksNodegroupStateSensor
18
19
# custom class to create a node group with Nodes on EKS
20
class EksCreateNodegroupWithNodesOperator(EksCreateNodegroupOperator):
21
    def execute(self, context):
22
        # instantiating an EKSHook on the basis of the AWS connection (Step 5)
23
        eks_hook = EksHook(
24
            aws_conn_id=self.aws_conn_id,
25
            region_name=self.region,
26
        )
27
28
        # define the Node group to create
29
        eks_hook.create_nodegroup(
30
            clusterName=self.cluster_name,
31
            nodegroupName=self.nodegroup_name,
32
            subnets=self.nodegroup_subnets,
33
            nodeRole=self.nodegroup_role_arn,
34
            scalingConfig={"minSize": 1, "maxSize": 1, "desiredSize": 1},
35
            diskSize=20,
36
            instanceTypes=["g4dn.xlarge"],
37
            amiType="AL2_x86_64_GPU",  # get GPU resources
38
            updateConfig={"maxUnavailable": 1},
39
        )
40
41
42
# instantiate the DAG
43
with DAG(
44
    start_date=datetime(2022, 6, 1),
45
    catchup=False,
46
    schedule="@daily",
47
    dag_id="KPO_remote_EKS_cluster_example_dag",
48
) as dag:
49
    # task 1 creates the node group
50
    create_gpu_nodegroup = EksCreateNodegroupWithNodesOperator(
51
        task_id="create_gpu_nodegroup",
52
        cluster_name="<your cluster name>",
53
        nodegroup_name="gpu-nodes",
54
        nodegroup_subnets=["<your subnet>", "<your subnet>"],
55
        nodegroup_role_arn="<arn of your EKS role>",
56
        aws_conn_id="<your aws conn id>",
57
        region="<your region>",
58
    )
59
60
    # task 2 check for node group status, if it is up and running
61
    check_nodegroup_status = EKSNodegroupStateSensor(
62
        task_id="check_nodegroup_status",
63
        cluster_name="<your cluster name>",
64
        nodegroup_name="gpu-nodes",
65
        mode="reschedule",
66
        timeout=60 * 30,
67
        exponential_backoff=True,
68
        aws_conn_id="<your aws conn id>",
69
        region="<your region>",
70
    )
71
72
    # task 3 the KubernetesPodOperator running a task
73
    # here, cluster_context and the kubernetes_conn_id are defined at the task level.
74
    run_on_EKS = KubernetesPodOperator(
75
        task_id="run_on_EKS",
76
        cluster_context="<arn of your cluster>",
77
        namespace="airflow-kpo-default",
78
        name="example_pod",
79
        image="ubuntu",
80
        cmds=["bash", "-cx"],
81
        arguments=["echo hello"],
82
        get_logs=True,
83
        in_cluster=False,
84
        kubernetes_conn_id="k8s",
85
        startup_timeout_seconds=240,
86
    )
87
88
    # task 4 deleting the node group
89
    delete_gpu_nodegroup = EksDeleteNodegroupOperator(
90
        task_id="delete_gpu_nodegroup",
91
        cluster_name="<your cluster name>",
92
        nodegroup_name="gpu-nodes",
93
        aws_conn_id="<your aws conn id>",
94
        region="<your region>",
95
    )
96
97
    # task 5 checking that the node group was deleted successfully
98
    check_nodegroup_termination = EksNodegroupStateSensor(
99
        task_id="check_nodegroup_termination",
100
        cluster_name="<your cluster name>",
101
        nodegroup_name="gpu-nodes",
102
        aws_conn_id="<your aws conn id>",
103
        region="<your region>",
104
        mode="reschedule",
105
        timeout=60 * 30,
106
        target_state=NodegroupStates.NONEXISTENT,
107
    )
108
109
    # setting the dependencies
110
    create_gpu_nodegroup >> check_nodegroup_status >> run_on_EKS
111
    run_on_EKS >> delete_gpu_nodegroup >> check_nodegroup_termination