ecs_run_task_flow_anonymized.py

from prefect import Task, task, Flow
from prefect.utilities.aws import get_boto_client
from prefect.utilities.tasks import defaults_from_attrs


class ECSRunTask(Task):
    def __init__(
        self,
        task_definition: str,
        cluster: str = "default",
        network_configuration: dict = None,
        boto_kwargs: dict = None,
        **kwargs
    ):
        self.task_definition = task_definition
        self.cluster = cluster
        self.network_configuration = network_configuration

        if boto_kwargs is None:
            self.boto_kwargs = {}
        else:
            self.boto_kwargs = boto_kwargs

        super().__init__(**kwargs)

    @defaults_from_attrs("task_definition", "cluster", "network_configuration")
    def run(
        self,
        task_definition: str = None,
        cluster: str = None,
        network_configuration: dict = None,
        credentials: dict = None,
    ):
        ecs_client = get_boto_client("ecs", credentials=credentials, **self.boto_kwargs)
        response = ecs_client.run_task(
            taskDefinition=task_definition,
            cluster=cluster,
            networkConfiguration=network_configuration,
        )
        return response


ecs = ECSRunTask(
    task_definition="your_task_def_family_name",
    cluster="your_cluster_name_or_arn",
    network_configuration={
        "awsvpcConfiguration": {
            "subnets": ["subnet-1", "subnet-2", "subnet-3"],
            "assignPublicIp": "ENABLED",  # if disabled, it won't be able to pull the image from ECR
        }
    },
)


@task(log_stdout=True)
def print_response(boto3_response: str):
    print(boto3_response)


with Flow("ecs_run_task") as flow:
    ecs_task = ecs()
    print_response(ecs_task)