Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment