Repository: apache/airflow

Issue:
Realtime ECS logging
**Description**

Currently when `ECSOperator` is run, the logs of the ECS task are fetched only when the task is done. That's not so convenient, especially when the task takes some good amount of time. In order to understand what's happening with the task, I need to go to Cloudwatch and search for the tasks logs. It would be good to have some parallel process that could fetch ECS task logs from Cloudwatch and make them visible in a realtime. 

**Are you willing to submit a PR?**

I can try, but I need to be guided.



Discussion:
That would be cool. There is a great example you can use - we have such streaming stackdriver task handler implemented already for GCS/stackdriver. It should be very similar.

 https://github.com/apache/airflow/blob/main/airflow/providers/google/cloud/log/stackdriver_task_handler.py 
@potiuk okay, you can assign to me, I will check what i can do
@potiuk i have checked the example and it looks like it's a bit different. From my understanding that task handler is used for logging for the operator. In the example writing of the logs can be sync/async depends on provided transport type. Reading is always synchronous. But anyway, in ECSOperator we run a separated ECS task with a separated log stream. And what we need is to read logs from ECS task log stream and write them to ECSOperator logs.

That's the code from ECSOperator:
```
        self._wait_for_task_ended()

        self._check_success_task()
```

`self._wait_for_task_ended()` is blocking. And we read logs of the ECS task in `self._check_success_task()`. So my question is how to take the part that reads the logs from `self._check_success_task()` and have it running in parallel with ` self._wait_for_task_ended()`

The only idea I have is not to use standard boto3 waiter, but have some custom wait implementation that pulls the status of the task and also pulls the logs in the same iteration. But that doesn't sound good.
Ah I see now.  I think you do not have to skip the waiter, it's quite likely you could be able to run separate thread reading the logs while the waiter waits.

Resolving patch:
diff --git a/airflow/providers/amazon/aws/operators/ecs.py b/airflow/providers/amazon/aws/operators/ecs.py
--- a/airflow/providers/amazon/aws/operators/ecs.py
+++ b/airflow/providers/amazon/aws/operators/ecs.py
@@ -17,10 +17,14 @@
 # under the License.
 import re
 import sys
+import time
 from collections import deque
-from datetime import datetime
+from datetime import datetime, timedelta
+from logging import Logger
+from threading import Event, Thread
 from typing import Dict, Generator, Optional
 
+from botocore.exceptions import ClientError
 from botocore.waiter import Waiter
 
 from airflow.exceptions import AirflowException
@@ -80,6 +84,73 @@ def list_tasks(self, cluster: str, launchType: str, desiredStatus: str, family:
         ...
 
 
+class ECSTaskLogFetcher(Thread):
+    """
+    Fetches Cloudwatch log events with specific interval as a thread
+    and sends the log events to the info channel of the provided logger.
+    """
+
+    def __init__(
+        self,
+        *,
+        aws_conn_id: Optional[str] = 'aws_default',
+        region_name: Optional[str] = None,
+        log_group: str,
+        log_stream_name: str,
+        fetch_interval: timedelta,
+        logger: Logger,
+    ):
+        super().__init__()
+        self._event = Event()
+
+        self.fetch_interval = fetch_interval
+
+        self.logger = logger
+        self.log_group = log_group
+        self.log_stream_name = log_stream_name
+
+        self.hook = AwsLogsHook(aws_conn_id=aws_conn_id, region_name=region_name)
+
+    def run(self) -> None:
+        logs_to_skip = 0
+        while not self.is_stopped():
+            log_events = self._get_log_events(logs_to_skip)
+            for log_event in log_events:
+                self.logger.info(self._event_to_str(log_event))
+                logs_to_skip += 1
+            time.sleep(self.fetch_interval.total_seconds())
+
+    def _get_log_events(self, skip: int = 0) -> Generator:
+        try:
+            yield from self.hook.get_log_events(self.log_group, self.log_stream_name, skip=skip)
+        except ClientError as error:
+            if error.response['Error']['Code'] != 'ResourceNotFoundException':
+                self.logger.warning('Error on retrieving Cloudwatch log events', error)
+
+            yield from ()
+
+    def _event_to_str(self, event: dict) -> str:
+        event_dt = datetime.utcfromtimestamp(event['timestamp'] / 1000.0)
+        formatted_event_dt = event_dt.strftime('%Y-%m-%d %H:%M:%S,%f')[:-3]
+        message = event['message']
+        return f'[{formatted_event_dt}] {message}'
+
+    def get_last_log_messages(self, number_messages) -> list:
+        return [log['message'] for log in deque(self._get_log_events(), maxlen=number_messages)]
+
+    def get_last_log_message(self) -> Optional[str]:
+        try:
+            return self.get_last_log_messages(1)[0]
+        except IndexError:
+            return None
+
+    def is_stopped(self) -> bool:
+        return self._event.is_set()
+
+    def stop(self):
+        self._event.set()
+
+
 class ECSOperator(BaseOperator):
     """
     Execute a task on AWS ECS (Elastic Container Service)
@@ -136,6 +207,9 @@ 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 awslogs_fetch_interval: the interval that the ECS task log fetcher should wait
+        in between each Cloudwatch logs fetches.
+    :type awslogs_fetch_interval: timedelta
     :param quota_retry: Config if and how to retry the launch of a new ECS task, to handle
         transient errors.
     :type quota_retry: dict
@@ -180,6 +254,7 @@ def __init__(
         awslogs_group: Optional[str] = None,
         awslogs_region: Optional[str] = None,
         awslogs_stream_prefix: Optional[str] = None,
+        awslogs_fetch_interval: timedelta = timedelta(seconds=30),
         propagate_tags: Optional[str] = None,
         quota_retry: Optional[dict] = None,
         reattach: bool = False,
@@ -205,6 +280,7 @@ def __init__(
         self.awslogs_group = awslogs_group
         self.awslogs_stream_prefix = awslogs_stream_prefix
         self.awslogs_region = awslogs_region
+        self.awslogs_fetch_interval = awslogs_fetch_interval
         self.propagate_tags = propagate_tags
         self.reattach = reattach
         self.number_logs_exception = number_logs_exception
@@ -216,6 +292,7 @@ def __init__(
         self.client: Optional[ECSProtocol] = None
         self.arn: Optional[str] = None
         self.retry_args = quota_retry
+        self.task_log_fetcher: Optional[ECSTaskLogFetcher] = None
 
     @provide_session
     def execute(self, context, session=None):
@@ -232,7 +309,19 @@ def execute(self, context, session=None):
         if not self.arn:
             self._start_task(context)
 
-        self._wait_for_task_ended()
+        if self._aws_logs_enabled():
+            self.log.info('Starting ECS Task Log Fetcher')
+            self.task_log_fetcher = self._get_task_log_fetcher()
+            self.task_log_fetcher.start()
+
+            try:
+                self._wait_for_task_ended()
+            finally:
+                self.task_log_fetcher.stop()
+
+            self.task_log_fetcher.join()
+        else:
+            self._wait_for_task_ended()
 
         self._check_success_task()
 
@@ -243,8 +332,8 @@ def execute(self, context, session=None):
             # 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()
+        if self.do_xcom_push and self.task_log_fetcher:
+            return self.task_log_fetcher.get_last_log_message()
 
         return None
 
@@ -286,8 +375,8 @@ def _start_task(self, context):
         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}")
+        self.ecs_task_id = self.arn.split("/")[-1]
+        self.log.info(f"ECS task ID is: {self.ecs_task_id}")
 
         if self.reattach:
             # Save the task ARN in XCom to be able to reattach it if needed
@@ -338,25 +427,20 @@ def _wait_for_task_ended(self) -> None:
 
         return
 
-    def _cloudwatch_log_events(self) -> Generator:
-        if self._aws_logs_enabled():
-            task_id = self.arn.split("/")[-1]
-            stream_name = f"{self.awslogs_stream_prefix}/{task_id}"
-            yield from self.get_logs_hook().get_log_events(self.awslogs_group, stream_name)
-        else:
-            yield from ()
-
     def _aws_logs_enabled(self):
         return self.awslogs_group and self.awslogs_stream_prefix
 
-    def _last_log_messages(self, number_messages):
-        return [log["message"] for log in deque(self._cloudwatch_log_events(), maxlen=number_messages)]
+    def _get_task_log_fetcher(self) -> ECSTaskLogFetcher:
+        log_stream_name = f"{self.awslogs_stream_prefix}/{self.ecs_task_id}"
 
-    def _last_log_message(self):
-        try:
-            return self._last_log_messages(1)[0]
-        except IndexError:
-            return None
+        return ECSTaskLogFetcher(
+            aws_conn_id=self.aws_conn_id,
+            region_name=self.awslogs_region,
+            log_group=self.awslogs_group,
+            log_stream_name=log_stream_name,
+            fetch_interval=self.awslogs_fetch_interval,
+            logger=self.log,
+        )
 
     def _check_success_task(self) -> None:
         if not self.client or not self.arn:
@@ -365,11 +449,6 @@ def _check_success_task(self) -> None:
         response = self.client.describe_tasks(cluster=self.cluster, tasks=[self.arn])
         self.log.info('ECS Task stopped, check status: %s', response)
 
-        # Get logs from CloudWatch if the awslogs log driver was used
-        for event in self._cloudwatch_log_events():
-            event_dt = datetime.fromtimestamp(event['timestamp'] / 1000.0)
-            self.log.info("[%s] %s", event_dt.isoformat(), event['message'])
-
         if len(response.get('failures', [])) > 0:
             raise AirflowException(response)
 
@@ -387,7 +466,9 @@ def _check_success_task(self) -> None:
             containers = task['containers']
             for container in containers:
                 if container.get('lastStatus') == 'STOPPED' and container['exitCode'] != 0:
-                    last_logs = "\n".join(self._last_log_messages(self.number_logs_exception))
+                    last_logs = "\n".join(
+                        self.task_log_fetcher.get_last_log_messages(self.number_logs_exception)
+                    )
                     raise AirflowException(
                         f"This task is not in success state - last {self.number_logs_exception} "
                         f"logs from Cloudwatch:\n{last_logs}"
@@ -409,14 +490,13 @@ def get_hook(self) -> AwsBaseHook:
         self.hook = AwsBaseHook(aws_conn_id=self.aws_conn_id, client_type='ecs', region_name=self.region_name)
         return self.hook
 
-    def get_logs_hook(self) -> AwsLogsHook:
-        """Create and return an AwsLogsHook."""
-        return AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.awslogs_region)
-
     def on_kill(self) -> None:
         if not self.client or not self.arn:
             return
 
+        if self.task_log_fetcher:
+            self.task_log_fetcher.stop()
+
         response = self.client.stop_task(
             cluster=self.cluster, task=self.arn, reason='Task killed by the user'
         )


Test patch:
diff --git a/tests/providers/amazon/aws/operators/test_ecs.py b/tests/providers/amazon/aws/operators/test_ecs.py
--- a/tests/providers/amazon/aws/operators/test_ecs.py
+++ b/tests/providers/amazon/aws/operators/test_ecs.py
@@ -20,14 +20,16 @@
 import sys
 import unittest
 from copy import deepcopy
+from datetime import timedelta
 from unittest import mock
 
 import pytest
+from botocore.exceptions import ClientError
 from parameterized import parameterized
 
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.exceptions import ECSOperatorError
-from airflow.providers.amazon.aws.operators.ecs import ECSOperator, should_retry
+from airflow.providers.amazon.aws.operators.ecs import ECSOperator, ECSTaskLogFetcher, should_retry
 
 # fmt: off
 RESPONSE_WITHOUT_FAILURES = {
@@ -244,26 +246,7 @@ def test_wait_end_tasks(self):
         client_mock.get_waiter.return_value.wait.assert_called_once_with(cluster='c', tasks=['arn'])
         assert sys.maxsize == client_mock.get_waiter.return_value.config.max_attempts
 
-    @mock.patch.object(
-        ECSOperator, '_cloudwatch_log_events', return_value=({"message": str(i)} for i in range(10))
-    )
-    def test_last_log_messages(self, mock_cloudwatch_log_events):
-        client_mock = mock.Mock()
-        self.ecs.arn = 'arn'
-        self.ecs.client = client_mock
-
-        assert self.ecs._last_log_messages(5) == ["5", "6", "7", "8", "9"]
-
-    @mock.patch.object(ECSOperator, '_cloudwatch_log_events', return_value=())
-    def test_last_log_messages_empty(self, mock_cloudwatch_log_events):
-        client_mock = mock.Mock()
-        self.ecs.arn = 'arn'
-        self.ecs.client = client_mock
-
-        assert self.ecs._last_log_messages(10) == []
-
-    @mock.patch.object(ECSOperator, '_last_log_messages', return_value=["1", "2", "3", "4", "5"])
-    def test_check_success_tasks_raises_cloudwatch_logs(self, mock_last_log_messages):
+    def test_check_success_tasks_raises_cloudwatch_logs(self):
         client_mock = mock.Mock()
         self.ecs.arn = 'arn'
         self.ecs.client = client_mock
@@ -271,6 +254,12 @@ def test_check_success_tasks_raises_cloudwatch_logs(self, mock_last_log_messages
         client_mock.describe_tasks.return_value = {
             'tasks': [{'containers': [{'name': 'foo', 'lastStatus': 'STOPPED', 'exitCode': 1}]}]
         }
+
+        task_log_fetcher = mock.Mock()
+        self.ecs.task_log_fetcher = task_log_fetcher
+
+        task_log_fetcher.get_last_log_messages.return_value = ["1", "2", "3", "4", "5"]
+
         with pytest.raises(Exception) as ctx:
             self.ecs._check_success_task()
 
@@ -279,8 +268,7 @@ def test_check_success_tasks_raises_cloudwatch_logs(self, mock_last_log_messages
         )
         client_mock.describe_tasks.assert_called_once_with(cluster='c', tasks=['arn'])
 
-    @mock.patch.object(ECSOperator, '_last_log_messages', return_value=[])
-    def test_check_success_tasks_raises_cloudwatch_logs_empty(self, mock_last_log_messages):
+    def test_check_success_tasks_raises_cloudwatch_logs_empty(self):
         client_mock = mock.Mock()
         self.ecs.arn = 'arn'
         self.ecs.client = client_mock
@@ -288,6 +276,12 @@ def test_check_success_tasks_raises_cloudwatch_logs_empty(self, mock_last_log_me
         client_mock.describe_tasks.return_value = {
             'tasks': [{'containers': [{'name': 'foo', 'lastStatus': 'STOPPED', 'exitCode': 1}]}]
         }
+
+        task_log_fetcher = mock.Mock()
+        self.ecs.task_log_fetcher = task_log_fetcher
+
+        task_log_fetcher.get_last_log_messages.return_value = []
+
         with pytest.raises(Exception) as ctx:
             self.ecs._check_success_task()
 
@@ -474,19 +468,26 @@ def test_reattach_save_task_arn_xcom(
         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_messages', return_value=["Log output"])
-    def test_execute_xcom_with_log(self, mock_cloudwatch_log_message):
+    def test_execute_xcom_with_log(self):
         self.ecs.do_xcom_push = True
+        self.ecs.task_log_fetcher = mock.Mock()
+        self.ecs.task_log_fetcher.get_last_log_message.return_value = "Log output"
         assert self.ecs.execute(None) == "Log output"
 
-    @mock.patch.object(ECSOperator, '_last_log_messages', return_value=[])
-    def test_execute_xcom_with_no_log(self, mock_cloudwatch_log_message):
+    def test_execute_xcom_with_no_log(self):
+        self.ecs.do_xcom_push = True
+        self.ecs.task_log_fetcher = mock.Mock()
+        self.ecs.task_log_fetcher.get_last_log_message.return_value = None
+        assert self.ecs.execute(None) is None
+
+    def test_execute_xcom_with_no_log_fetcher(self):
         self.ecs.do_xcom_push = True
         assert self.ecs.execute(None) is None
 
-    @mock.patch.object(ECSOperator, '_last_log_messages', return_value=["Log output"])
-    def test_execute_xcom_disabled(self, mock_cloudwatch_log_message):
+    def test_execute_xcom_disabled(self):
         self.ecs.do_xcom_push = False
+        self.ecs.task_log_fetcher = mock.Mock()
+        self.ecs.task_log_fetcher.get_last_log_message.return_value = "Log output"
         assert self.ecs.execute(None) is None
 
 
@@ -496,3 +497,121 @@ def test_return_true_on_valid_reason(self):
 
     def test_return_false_on_invalid_reason(self):
         self.assertFalse(should_retry(ECSOperatorError([{'reason': 'CLUSTER_NOT_FOUND'}], 'Foo')))
+
+
+class TestECSTaskLogFetcher(unittest.TestCase):
+    @mock.patch('logging.Logger')
+    def set_up_log_fetcher(self, logger_mock):
+        self.logger_mock = logger_mock
+
+        self.log_fetcher = ECSTaskLogFetcher(
+            log_group="test_log_group",
+            log_stream_name="test_log_stream_name",
+            fetch_interval=timedelta(milliseconds=1),
+            logger=logger_mock,
+        )
+
+    def setUp(self):
+        self.set_up_log_fetcher()
+
+    @mock.patch(
+        'threading.Event.is_set',
+        side_effect=(False, False, False, True),
+    )
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        side_effect=(
+            iter(
+                [
+                    {'timestamp': 1617400267123, 'message': 'First'},
+                    {'timestamp': 1617400367456, 'message': 'Second'},
+                ]
+            ),
+            iter(
+                [
+                    {'timestamp': 1617400467789, 'message': 'Third'},
+                ]
+            ),
+            iter([]),
+        ),
+    )
+    def test_run(self, get_log_events_mock, event_is_set_mock):
+
+        self.log_fetcher.run()
+
+        self.logger_mock.info.assert_has_calls(
+            [
+                mock.call('[2021-04-02 21:51:07,123] First'),
+                mock.call('[2021-04-02 21:52:47,456] Second'),
+                mock.call('[2021-04-02 21:54:27,789] Third'),
+            ]
+        )
+
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        side_effect=ClientError({"Error": {"Code": "ResourceNotFoundException"}}, None),
+    )
+    def test_get_log_events_with_expected_error(self, get_log_events_mock):
+        with pytest.raises(StopIteration):
+            next(self.log_fetcher._get_log_events())
+
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        side_effect=Exception(),
+    )
+    def test_get_log_events_with_unexpected_error(self, get_log_events_mock):
+        with pytest.raises(Exception):
+            next(self.log_fetcher._get_log_events())
+
+    def test_event_to_str(self):
+        events = [
+            {'timestamp': 1617400267123, 'message': 'First'},
+            {'timestamp': 1617400367456, 'message': 'Second'},
+            {'timestamp': 1617400467789, 'message': 'Third'},
+        ]
+        assert [self.log_fetcher._event_to_str(event) for event in events] == (
+            [
+                '[2021-04-02 21:51:07,123] First',
+                '[2021-04-02 21:52:47,456] Second',
+                '[2021-04-02 21:54:27,789] Third',
+            ]
+        )
+
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        return_value=(),
+    )
+    def test_get_last_log_message_with_no_log_events(self, mock_log_events):
+        assert self.log_fetcher.get_last_log_message() is None
+
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        return_value=iter(
+            [
+                {'timestamp': 1617400267123, 'message': 'First'},
+                {'timestamp': 1617400367456, 'message': 'Second'},
+            ]
+        ),
+    )
+    def test_get_last_log_message_with_log_events(self, mock_log_events):
+        assert self.log_fetcher.get_last_log_message() == 'Second'
+
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        return_value=iter(
+            [
+                {'timestamp': 1617400267123, 'message': 'First'},
+                {'timestamp': 1617400367456, 'message': 'Second'},
+                {'timestamp': 1617400367458, 'message': 'Third'},
+            ]
+        ),
+    )
+    def test_get_last_log_messages_with_log_events(self, mock_log_events):
+        assert self.log_fetcher.get_last_log_messages(2) == ['Second', 'Third']
+
+    @mock.patch(
+        'airflow.providers.amazon.aws.hooks.logs.AwsLogsHook.get_log_events',
+        return_value=(),
+    )
+    def test_get_last_log_messages_with_no_log_events(self, mock_log_events):
+        assert self.log_fetcher.get_last_log_messages(2) == []