ecsruntask.py

from prefect import Task
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)
        # todo: if network configuration is none, try to infer that as in ECS agent
        response = ecs_client.run_task(
            taskDefinition=task_definition,
            cluster=cluster,
            networkConfiguration=network_configuration,
        )

        return response