Various fixes on ECS run task operator

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

@@ -185,7 +185,14 @@ def _get_log_events(self, skip: int = 0) -> Generator:
  except ClientError as error:
  if error.response["Error"]["Code"] != "ResourceNotFoundException":
  self.logger.warning("Error on retrieving Cloudwatch log events", error)
-
+ else:
+ self.logger.info(
+ "Cannot find log stream yet, it can take a couple of seconds to show up. "
+ "If this error persists, check that the log group and stream are correct: "
+ "group: %s\tstream: %s",
+ self.log_group,
+ self.log_stream_name,
+ )
  yield from ()
  except ConnectionClosedError as error:
  self.logger.warning("ConnectionClosedError on retrieving Cloudwatch log events", error)

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

@@ -480,6 +480,17 @@ def __init__(
  self.waiter_delay = waiter_delay
  self.waiter_max_attempts = waiter_max_attempts
 
+ if self._aws_logs_enabled() and not self.wait_for_completion:
+ self.log.warning(
+ "Trying to get logs without waiting for the task to complete is undefined behavior."
+ )
+
+ @staticmethod
+ def _get_ecs_task_id(task_arn: str | None) -> str | None:
+ if task_arn is None:
+ return None
+ return task_arn.split("/")[-1]
+
  @provide_session
  def execute(self, context, session=None):
  self.log.info(
@@ -506,25 +517,24 @@ def execute(self, context, session=None):
 
  @AwsBaseHook.retry(should_retry_eni)
  def _start_wait_check_task(self, context):
-
  if not self.arn:
  self._start_task(context)
 
+ if not self.wait_for_completion:
+ return
+
  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:
- if self.wait_for_completion:
- self._wait_for_task_ended()
+ self._wait_for_task_ended()
  finally:
  self.task_log_fetcher.stop()
-
  self.task_log_fetcher.join()
  else:
- if self.wait_for_completion:
- self._wait_for_task_ended()
+ self._wait_for_task_ended()
 
  self._check_success_task()
 
@@ -566,8 +576,7 @@ def _start_task(self, context):
  self.log.info("ECS Task started: %s", response)
 
  self.arn = response["tasks"][0]["taskArn"]
- self.ecs_task_id = self.arn.split("/")[-1]
- self.log.info("ECS task ID is: %s", self.ecs_task_id)
+ self.log.info("ECS task ID is: %s", self._get_ecs_task_id(self.arn))
 
  if self.reattach:
  # Save the task ARN in XCom to be able to reattach it if needed
@@ -590,7 +599,6 @@ def _try_reattach_task(self, context):
  )
  if previous_task_arn in running_tasks:
  self.arn = previous_task_arn
- self.ecs_task_id = self.arn.split("/")[-1]
  self.log.info("Reattaching previously launched task: %s", self.arn)
  else:
  self.log.info("No active previously launched task found to reattach")
@@ -620,7 +628,7 @@ def _aws_logs_enabled(self):
  def _get_task_log_fetcher(self) -> EcsTaskLogFetcher:
  if not self.awslogs_group:
  raise ValueError("must specify awslogs_group to fetch task logs")
- log_stream_name = f"{self.awslogs_stream_prefix}/{self.ecs_task_id}"
+ log_stream_name = f"{self.awslogs_stream_prefix}/{self._get_ecs_task_id(self.arn)}"
 
  return EcsTaskLogFetcher(
  aws_conn_id=self.aws_conn_id,

docs/apache-airflow-providers-amazon/operators/ecs.rst

@@ -250,7 +250,7 @@ both can be overridden with provided values. Raises an AirflowException with
 the failure reason if a failed state is provided and that state is reached
 before the target state.
 
-.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_ecs.py
+.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_ecs_fargate.py
  :language: python
  :dedent: 4
  :start-after: [START howto_sensor_ecs_task_state]

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

@@ -304,7 +304,10 @@ def test_execute_without_failures(
  wait_mock.assert_called_once_with()
  check_mock.assert_called_once_with()
  assert self.ecs.arn == f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}"
- assert self.ecs.ecs_task_id == TASK_ID
+
+ def test_task_id_parsing(self):
+ id = EcsRunTaskOperator._get_ecs_task_id(f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}")
+ assert id == TASK_ID
 
  @mock.patch.object(EcsBaseOperator, "client")
  def test_execute_with_failures(self, client_mock):
@@ -571,7 +574,6 @@ def test_reattach_successful(
  check_mock.assert_called_once_with()
  xcom_del_mock.assert_called_once()
  assert self.ecs.arn == f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}"
- assert self.ecs.ecs_task_id == TASK_ID
 
  @pytest.mark.parametrize(
  "launch_type, tags",
@@ -620,7 +622,6 @@ def test_reattach_save_task_arn_xcom(
  check_mock.assert_called_once_with()
  xcom_del_mock.assert_called_once()
  assert self.ecs.arn == f"arn:aws:ecs:us-east-1:012345678910:task/{TASK_ID}"
- assert self.ecs.ecs_task_id == TASK_ID
 
  @mock.patch.object(EcsBaseOperator, "client")
  @mock.patch("airflow.providers.amazon.aws.hooks.ecs.EcsTaskLogFetcher")

tests/system/providers/amazon/aws/example_ecs.py

@@ -23,7 +23,7 @@
 from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
-from airflow.providers.amazon.aws.hooks.ecs import EcsClusterStates, EcsTaskStates
+from airflow.providers.amazon.aws.hooks.ecs import EcsClusterStates
 from airflow.providers.amazon.aws.operators.ecs import (
  EcsCreateClusterOperator,
  EcsDeleteClusterOperator,
@@ -34,7 +34,6 @@
 from airflow.providers.amazon.aws.sensors.ecs import (
  EcsClusterStateSensor,
  EcsTaskDefinitionStateSensor,
- EcsTaskStateSensor,
 )
 from airflow.utils.trigger_rule import TriggerRule
 from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
@@ -67,6 +66,15 @@ def get_region():
  return boto3.session.Session().region_name
 
 
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def clean_logs(group_name: str):
+ client = boto3.client("logs")
+ # A bit brutal to delete the whole group, I know,
+ # but we don't have the access to the arn of the task which is used in the stream name
+ # and also those logs just contain "hello world", which is not very interesting.
+ client.delete_log_group(logGroupName=group_name)
+
+
 with DAG(
  dag_id=DAG_ID,
  schedule="@once",
@@ -85,6 +93,7 @@ def get_region():
  asg_name = f"{env_id}-asg"
 
  aws_region = get_region()
+ log_group_name = "/ecs/hello-world"
 
  # [START howto_operator_ecs_create_cluster]
  create_cluster = EcsCreateClusterOperator(
@@ -114,7 +123,16 @@ def get_region():
  "workingDirectory": "/usr/bin",
  "entryPoint": ["sh", "-c"],
  "command": ["ls"],
- }
+ "logConfiguration": {
+ "logDriver": "awslogs",
+ "options": {
+ "awslogs-group": log_group_name,
+ "awslogs-region": aws_region,
+ "awslogs-create-group": "true",
+ "awslogs-stream-prefix": "ecs",
+ },
+ },
+ },
  ],
  register_task_kwargs={
  "cpu": "256",
@@ -140,38 +158,19 @@ def get_region():
  "containerOverrides": [
  {
  "name": container_name,
- "command": ["echo", "hello", "world"],
+ "command": ["echo hello world"],
  },
  ],
  },
  network_configuration={"awsvpcConfiguration": {"subnets": existing_cluster_subnets}},
  # [START howto_awslogs_ecs]
- awslogs_group="/ecs/hello-world",
+ awslogs_group=log_group_name,
  awslogs_region=aws_region,
- awslogs_stream_prefix="ecs/hello-world-container",
+ awslogs_stream_prefix=f"ecs/{container_name}",
  # [END howto_awslogs_ecs]
- # You must set `reattach=True` in order to get ecs_task_arn if you plan to use a Sensor.
- reattach=True,
  )
  # [END howto_operator_ecs_run_task]
 
- # EcsRunTaskOperator waits by default, setting as False to test the Sensor below.
- run_task.wait_for_completion = False
-
- # [START howto_sensor_ecs_task_state]
- # By default, EcsTaskStateSensor waits until the task has started, but the
- # demo task runs so fast that the sensor misses it. This sensor instead
- # demonstrates how to wait until the ECS Task has completed by providing
- # the target_state and failure_states parameters.
- await_task_finish = EcsTaskStateSensor(
- task_id="await_task_finish",
- cluster=existing_cluster_name,
- task=run_task.output["ecs_task_arn"],
- target_state=EcsTaskStates.STOPPED,
- failure_states={EcsTaskStates.NONE},
- )
- # [END howto_sensor_ecs_task_state]
-
  # [START howto_operator_ecs_deregister_task_definition]
  deregister_task = EcsDeregisterTaskDefinitionOperator(
  task_id="deregister_task",
@@ -209,10 +208,10 @@ def get_region():
  register_task,
  await_task_definition,
  run_task,
- await_task_finish,
  deregister_task,
  delete_cluster,
  await_delete_cluster,
+ clean_logs(log_group_name),
  )
 
  from tests.system.utils.watcher import watcher

tests/system/providers/amazon/aws/example_ecs_fargate.py

@@ -23,7 +23,9 @@
 from airflow import DAG
 from airflow.decorators import task
 from airflow.models.baseoperator import chain
+from airflow.providers.amazon.aws.hooks.ecs import EcsTaskStates
 from airflow.providers.amazon.aws.operators.ecs import EcsRunTaskOperator
+from airflow.providers.amazon.aws.sensors.ecs import EcsTaskStateSensor
 from airflow.utils.trigger_rule import TriggerRule
 from tests.system.providers.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
 
@@ -120,16 +122,36 @@ def delete_cluster(cluster_name: str) -> None:
  "assignPublicIp": "ENABLED",
  },
  },
+ # You must set `reattach=True` in order to get ecs_task_arn if you plan to use a Sensor.
+ reattach=True,
  )
  # [END howto_operator_ecs]
 
+ # EcsRunTaskOperator waits by default, setting as False to test the Sensor below.
+ hello_world.wait_for_completion = False
+
+ # [START howto_sensor_ecs_task_state]
+ # By default, EcsTaskStateSensor waits until the task has started, but the
+ # demo task runs so fast that the sensor misses it. This sensor instead
+ # demonstrates how to wait until the ECS Task has completed by providing
+ # the target_state and failure_states parameters.
+ await_task_finish = EcsTaskStateSensor(
+ task_id="await_task_finish",
+ cluster=cluster_name,
+ task=hello_world.output["ecs_task_arn"],
+ target_state=EcsTaskStates.STOPPED,
+ failure_states={EcsTaskStates.NONE},
+ )
+ # [END howto_sensor_ecs_task_state]
+
  chain(
  # TEST SETUP
  test_context,
  create_cluster(cluster_name),
  create_task_definition,
  # TEST BODY
  hello_world,
+ await_task_finish,
  # TEST TEARDOWN
  delete_task_definition(create_task_definition),
  delete_cluster(cluster_name),