Allow attaching to previously launched task in ECSOperator (#16685)

This PR adds a parameter reattach_prev_task to ECSOperator.

**Before:**
Until now we could use 'reattach' which was reattaching a running ECS Task (if there was one running) of the same 'family' instead of creating a new one. The problem was that if we had workflows using the same ECS Task Definition in several tasks, it didn't know which one to reattach and we could only use `concurrency=1` in some pipelines for example (when we launch the same ECS task in parallel from Airflow with different configurations).

**Now:**
Now with reattach_prev_task instead, when we launch a new ECS task, it will store temporarily the ECS Task ARN in XCOM. If there is an issue during the run (typically connection problem between Airflow and ECS for long-running tasks or Airflow worker restarting which was then still running those tasks in the background without Airflow being aware of it):
- self._start_task will store the ECS task ARN in XCOM (in a 'fake' task_id equal to f"{self.task_id}_task_arn"
- in the next execution, it will check if this task ARN is still running and if so it will reattach it to the operator, otherwise it will create a new one
- when the operator runs succesfully it will delete the XCOM value


I didn't change the logic of 'reattach' to do that directly because I didn't know if it had been designed for other use cases

**Update 2021-07-01:**
After discussing with @darwinyip  I made the change to 'reattach' directly instead of creating a new flag

Author: pmalafosse

airflow/providers/amazon/aws/hooks/base_aws.py

@@ -530,7 +530,7 @@ def retry_decorator(fun: Callable):
             def decorator_f(self, *args, **kwargs):
                 retry_args = getattr(self, 'retry_args', None)
                 if retry_args is None:
-                    return fun(self)
+                    return fun(self, *args, **kwargs)
                 multiplier = retry_args.get('multiplier', 1)
                 min_limit = retry_args.get('min', 1)
                 max_limit = retry_args.get('max', 1)
@@ -543,7 +543,7 @@ def decorator_f(self, *args, **kwargs):
                     'before': tenacity_logger,
                     'after': tenacity_logger,
                 }
-                return tenacity.retry(**default_kwargs)(fun)(self)
+                return tenacity.retry(**default_kwargs)(fun)(self, *args, **kwargs)
 
             return decorator_f
 

airflow/providers/amazon/aws/operators/ecs.py

@@ -24,11 +24,12 @@
 from botocore.waiter import Waiter
 
 from airflow.exceptions import AirflowException
-from airflow.models import BaseOperator
+from airflow.models import BaseOperator, XCom
 from airflow.providers.amazon.aws.exceptions import ECSOperatorError
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.typing_compat import Protocol, runtime_checkable
+from airflow.utils.session import provide_session
 
 
 def should_retry(exception: Exception):
@@ -135,8 +136,10 @@ class ECSOperator(BaseOperator):
         Only required if you want logs to be shown in the Airflow UI after your job has
         finished.
     :type awslogs_stream_prefix: str
-    :param reattach: If set to True, will check if a task from the same family is already running.
-        If so, the operator will attach to it instead of starting a new task.
+    :param reattach: If set to True, will check if the task previously launched by the task_instance
+        is already running. If so, the operator will attach to it instead of starting a new task.
+        This is to avoid relaunching a new task when the connection drops between Airflow and ECS while
+        the task is running (when the Airflow worker is restarted for example).
     :type reattach: bool
     :param quota_retry: Config if and how to retry _start_task() for transient errors.
     :type quota_retry: dict
@@ -145,6 +148,8 @@ class ECSOperator(BaseOperator):
     ui_color = '#f0ede4'
     template_fields = ('overrides',)
     template_fields_renderers = {"overrides": "json"}
+    REATTACH_XCOM_KEY = "ecs_task_arn"
+    REATTACH_XCOM_TASK_ID_TEMPLATE = "{task_id}_task_arn"
 
     def __init__(
         self,
@@ -200,7 +205,8 @@ def __init__(
         self.arn: Optional[str] = None
         self.retry_args = quota_retry
 
-    def execute(self, context):
+    @provide_session
+    def execute(self, context, session=None):
         self.log.info(
             'Running ECS Task - Task definition: %s - on cluster %s', self.task_definition, self.cluster
         )
@@ -212,20 +218,28 @@ def execute(self, context):
             self._try_reattach_task()
 
         if not self.arn:
-            self._start_task()
+            self._start_task(context)
 
         self._wait_for_task_ended()
 
         self._check_success_task()
 
         self.log.info('ECS Task has been successfully executed')
 
+        if self.reattach:
+            # Clear the XCom value storing the ECS task ARN if the task has completed
+            # as we can't reattach it anymore
+            self._xcom_del(session, self.REATTACH_XCOM_TASK_ID_TEMPLATE.format(task_id=self.task_id))
+
         if self.do_xcom_push:
             return self._last_log_message()
 
         return None
 
-    def _start_task(self):
+    def _xcom_del(self, session, task_id):
+        session.query(XCom).filter(XCom.dag_id == self.dag_id, XCom.task_id == task_id).delete()
+
+    def _start_task(self, context):
         run_opts = {
             'cluster': self.cluster,
             'taskDefinition': self.task_definition,
@@ -261,6 +275,26 @@ def _start_task(self):
         self.log.info('ECS Task started: %s', response)
 
         self.arn = response['tasks'][0]['taskArn']
+        ecs_task_id = self.arn.split("/")[-1]
+        self.log.info(f"ECS task ID is: {ecs_task_id}")
+
+        if self.reattach:
+            # Save the task ARN in XCom to be able to reattach it if needed
+            self._xcom_set(
+                context,
+                key=self.REATTACH_XCOM_KEY,
+                value=self.arn,
+                task_id=self.REATTACH_XCOM_TASK_ID_TEMPLATE.format(task_id=self.task_id),
+            )
+
+    def _xcom_set(self, context, key, value, task_id):
+        XCom.set(
+            key=key,
+            value=value,
+            task_id=task_id,
+            dag_id=self.dag_id,
+            execution_date=context["ti"].execution_date,
+        )
 
     def _try_reattach_task(self):
         task_def_resp = self.client.describe_task_definition(taskDefinition=self.task_definition)
@@ -271,15 +305,16 @@ def _try_reattach_task(self):
         )
         running_tasks = list_tasks_resp['taskArns']
 
-        running_tasks_count = len(running_tasks)
-        if running_tasks_count > 1:
-            self.arn = running_tasks[0]
-            self.log.warning('More than 1 ECS Task found. Reattaching to %s', self.arn)
-        elif running_tasks_count == 1:
-            self.arn = running_tasks[0]
-            self.log.info('Reattaching task: %s', self.arn)
+        # Check if the ECS task previously launched is already running
+        previous_task_arn = self.xcom_pull(
+            task_ids=self.REATTACH_XCOM_TASK_ID_TEMPLATE.format(task_id=self.task_id),
+            key=self.REATTACH_XCOM_KEY,
+        )
+        if previous_task_arn in running_tasks:
+            self.arn = previous_task_arn
+            self.log.info("Reattaching previously launched task: %s", self.arn)
         else:
-            self.log.info('No active tasks found to reattach')
+            self.log.info("No active previously launched task found to reattach")
 
     def _wait_for_task_ended(self) -> None:
         if not self.client or not self.arn:

tests/providers/amazon/aws/operators/test_ecs.py

@@ -316,16 +316,27 @@ def test_check_success_task_not_raises(self):
             ['', {'testTagKey': 'testTagValue'}],
         ]
     )
+    @mock.patch.object(ECSOperator, "_xcom_del")
+    @mock.patch.object(
+        ECSOperator,
+        "xcom_pull",
+        return_value="arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55",
+    )
     @mock.patch.object(ECSOperator, '_wait_for_task_ended')
     @mock.patch.object(ECSOperator, '_check_success_task')
     @mock.patch.object(ECSOperator, '_start_task')
-    def test_reattach_successful(self, launch_type, tags, start_mock, check_mock, wait_mock):
+    def test_reattach_successful(
+        self, launch_type, tags, start_mock, check_mock, wait_mock, xcom_pull_mock, xcom_del_mock
+    ):
 
-        self.set_up_operator(launch_type=launch_type, tags=tags)
+        self.set_up_operator(launch_type=launch_type, tags=tags)  # pylint: disable=no-value-for-parameter
         client_mock = self.aws_hook_mock.return_value.get_conn.return_value
         client_mock.describe_task_definition.return_value = {'taskDefinition': {'family': 'f'}}
         client_mock.list_tasks.return_value = {
-            'taskArns': ['arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55']
+            'taskArns': [
+                'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b54',
+                'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55',
+            ]
         }
 
         self.ecs.reattach = True
@@ -345,8 +356,61 @@ def test_reattach_successful(self, launch_type, tags, start_mock, check_mock, wa
         client_mock.list_tasks.assert_called_once_with(cluster='c', desiredStatus='RUNNING', family='f')
 
         start_mock.assert_not_called()
+        xcom_pull_mock.assert_called_once_with(
+            key=self.ecs.REATTACH_XCOM_KEY,
+            task_ids=self.ecs.REATTACH_XCOM_TASK_ID_TEMPLATE.format(task_id=self.ecs.task_id),
+        )
+        wait_mock.assert_called_once_with()
+        check_mock.assert_called_once_with()
+        xcom_del_mock.assert_called_once()
+        assert self.ecs.arn == 'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55'
+
+    @parameterized.expand(
+        [
+            ['EC2', None],
+            ['FARGATE', None],
+            ['EC2', {'testTagKey': 'testTagValue'}],
+            ['', {'testTagKey': 'testTagValue'}],
+        ]
+    )
+    @mock.patch.object(ECSOperator, '_xcom_del')
+    @mock.patch.object(ECSOperator, '_xcom_set')
+    @mock.patch.object(ECSOperator, '_try_reattach_task')
+    @mock.patch.object(ECSOperator, '_wait_for_task_ended')
+    @mock.patch.object(ECSOperator, '_check_success_task')
+    def test_reattach_save_task_arn_xcom(
+        self, launch_type, tags, check_mock, wait_mock, reattach_mock, xcom_set_mock, xcom_del_mock
+    ):
+
+        self.set_up_operator(launch_type=launch_type, tags=tags)  # pylint: disable=no-value-for-parameter
+        client_mock = self.aws_hook_mock.return_value.get_conn.return_value
+        client_mock.describe_task_definition.return_value = {'taskDefinition': {'family': 'f'}}
+        client_mock.list_tasks.return_value = {'taskArns': []}
+        client_mock.run_task.return_value = RESPONSE_WITHOUT_FAILURES
+
+        self.ecs.reattach = True
+        self.ecs.execute(None)
+
+        self.aws_hook_mock.return_value.get_conn.assert_called_once()
+        extend_args = {}
+        if launch_type:
+            extend_args['launchType'] = launch_type
+        if launch_type == 'FARGATE':
+            extend_args['platformVersion'] = 'LATEST'
+        if tags:
+            extend_args['tags'] = [{'key': k, 'value': v} for (k, v) in tags.items()]
+
+        reattach_mock.assert_called_once()
+        client_mock.run_task.assert_called_once()
+        xcom_set_mock.assert_called_once_with(
+            None,
+            key=self.ecs.REATTACH_XCOM_KEY,
+            task_id=self.ecs.REATTACH_XCOM_TASK_ID_TEMPLATE.format(task_id=self.ecs.task_id),
+            value="arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55",
+        )
         wait_mock.assert_called_once_with()
         check_mock.assert_called_once_with()
+        xcom_del_mock.assert_called_once()
         assert self.ecs.arn == 'arn:aws:ecs:us-east-1:012345678910:task/d8c67b3c-ac87-4ffe-a847-4785bc3a8b55'
 
     @mock.patch.object(ECSOperator, '_last_log_message', return_value="Log output")