apache-airflow-providers-amazon 8.17.0rc2__py3-none-any.whl → 8.18.0rc1__py3-none-any.whl

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

@@ -17,10 +17,10 @@
 from __future__ import annotations
 
 import asyncio
-import warnings
 from typing import Any, Sequence
 
 import botocore.exceptions
+from deprecated import deprecated
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseAsyncHook, AwsBaseHook
@@ -195,16 +195,17 @@ class RedshiftHook(AwsBaseHook):
             return None
 
 
+@deprecated(
+    reason=(
+        "airflow.providers.amazon.aws.hook.base_aws.RedshiftAsyncHook "
+        "has been deprecated and will be removed in future"
+    ),
+    category=AirflowProviderDeprecationWarning,
+)
 class RedshiftAsyncHook(AwsBaseAsyncHook):
     """Interact with AWS Redshift using aiobotocore library."""
 
     def __init__(self, *args, **kwargs):
-        warnings.warn(
-            "airflow.providers.amazon.aws.hook.base_aws.RedshiftAsyncHook has been deprecated and "
-            "will be removed in future",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
         kwargs["client_type"] = "redshift"
         super().__init__(*args, **kwargs)
 

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

@@ -200,7 +200,7 @@ class RedshiftSQLHook(DbApiHook):
         return redshift_connector.connect(**conn_kwargs)
 
     def get_openlineage_database_info(self, connection: Connection) -> DatabaseInfo:
-        """Returns Redshift specific information for OpenLineage."""
+        """Return Redshift specific information for OpenLineage."""
         from airflow.providers.openlineage.sqlparser import DatabaseInfo
 
         authority = self._get_openlineage_redshift_authority_part(connection)
@@ -252,9 +252,9 @@ class RedshiftSQLHook(DbApiHook):
             return hostname
 
     def get_openlineage_database_dialect(self, connection: Connection) -> str:
-        """Returns redshift dialect."""
+        """Return redshift dialect."""
         return "redshift"
 
     def get_openlineage_default_schema(self) -> str | None:
-        """Returns current schema. This is usually changed with ``SEARCH_PATH`` parameter."""
+        """Return current schema. This is usually changed with ``SEARCH_PATH`` parameter."""
         return self.get_first("SELECT CURRENT_SCHEMA();")[0]

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

@@ -1369,6 +1369,10 @@ class S3Hook(AwsBaseHook):
         """
         Download a file from the S3 location to the local file system.
 
+        Note:
+            This function shadows the 'download_file' method of S3 API, but it is not the same.
+            If you want to use the original method from S3 API, please use 'S3Hook.get_conn().download_file()'
+
         .. seealso::
             - :external+boto3:py:meth:`S3.Object.download_fileobj`
 
@@ -1386,12 +1390,6 @@ class S3Hook(AwsBaseHook):
             Default: True.
         :return: the file name.
         """
-        self.log.info(
-            "This function shadows the 'download_file' method of S3 API, but it is not the same. If you "
-            "want to use the original method from S3 API, please call "
-            "'S3Hook.get_conn().download_file()'"
-        )
-
         self.log.info("Downloading source S3 file from Bucket %s with path %s", bucket_name, key)
 
         try:

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

@@ -26,8 +26,9 @@ import warnings
 from collections import Counter, namedtuple
 from datetime import datetime
 from functools import partial
-from typing import Any, Callable, Generator, cast
+from typing import Any, AsyncGenerator, Callable, Generator, cast
 
+from asgiref.sync import sync_to_async
 from botocore.exceptions import ClientError
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
@@ -310,10 +311,12 @@ class SageMakerHook(AwsBaseHook):
                 max_ingestion_time,
             )
 
-            billable_time = (
-                describe_response["TrainingEndTime"] - describe_response["TrainingStartTime"]
-            ) * describe_response["ResourceConfig"]["InstanceCount"]
-            self.log.info("Billable seconds: %d", int(billable_time.total_seconds()) + 1)
+            billable_seconds = SageMakerHook.count_billable_seconds(
+                training_start_time=describe_response["TrainingStartTime"],
+                training_end_time=describe_response["TrainingEndTime"],
+                instance_count=describe_response["ResourceConfig"]["InstanceCount"],
+            )
+            self.log.info("Billable seconds: %d", billable_seconds)
 
         return response
 
@@ -811,10 +814,12 @@ class SageMakerHook(AwsBaseHook):
             if status in failed_states:
                 reason = last_description.get("FailureReason", "(No reason provided)")
                 raise AirflowException(f"Error training {job_name}: {status} Reason: {reason}")
-            billable_time = (
-                last_description["TrainingEndTime"] - last_description["TrainingStartTime"]
-            ) * instance_count
-            self.log.info("Billable seconds: %d", int(billable_time.total_seconds()) + 1)
+            billable_seconds = SageMakerHook.count_billable_seconds(
+                training_start_time=last_description["TrainingStartTime"],
+                training_end_time=last_description["TrainingEndTime"],
+                instance_count=instance_count,
+            )
+            self.log.info("Billable seconds: %d", billable_seconds)
 
     def list_training_jobs(
         self, name_contains: str | None = None, max_results: int | None = None, **kwargs
@@ -1300,3 +1305,125 @@ class SageMakerHook(AwsBaseHook):
             if "BestCandidate" in res:
                 return res["BestCandidate"]
         return None
+
+    @staticmethod
+    def count_billable_seconds(
+        training_start_time: datetime, training_end_time: datetime, instance_count: int
+    ) -> int:
+        billable_time = (training_end_time - training_start_time) * instance_count
+        return int(billable_time.total_seconds()) + 1
+
+    async def describe_training_job_async(self, job_name: str) -> dict[str, Any]:
+        """
+        Return the training job info associated with the name.
+
+        :param job_name: the name of the training job
+        """
+        async with self.async_conn as client:
+            response: dict[str, Any] = await client.describe_training_job(TrainingJobName=job_name)
+            return response
+
+    async def describe_training_job_with_log_async(
+        self,
+        job_name: str,
+        positions: dict[str, Any],
+        stream_names: list[str],
+        instance_count: int,
+        state: int,
+        last_description: dict[str, Any],
+        last_describe_job_call: float,
+    ) -> tuple[int, dict[str, Any], float]:
+        """
+        Return the training job info associated with job_name and print CloudWatch logs.
+
+        :param job_name: name of the job to check status
+        :param positions: A list of pairs of (timestamp, skip) which represents the last record
+            read from each stream.
+        :param stream_names: A list of the log stream names. The position of the stream in this list is
+            the stream number.
+        :param instance_count: Count of the instance created for the job initially
+        :param state: log state
+        :param last_description: Latest description of the training job
+        :param last_describe_job_call: previous job called time
+        """
+        log_group = "/aws/sagemaker/TrainingJobs"
+
+        if len(stream_names) < instance_count:
+            logs_hook = AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+            streams = await logs_hook.describe_log_streams_async(
+                log_group=log_group,
+                stream_prefix=job_name + "/",
+                order_by="LogStreamName",
+                count=instance_count,
+            )
+
+            stream_names = [s["logStreamName"] for s in streams["logStreams"]] if streams else []
+            positions.update([(s, Position(timestamp=0, skip=0)) for s in stream_names if s not in positions])
+
+        if len(stream_names) > 0:
+            async for idx, event in self.get_multi_stream(log_group, stream_names, positions):
+                self.log.info(event["message"])
+                ts, count = positions[stream_names[idx]]
+                if event["timestamp"] == ts:
+                    positions[stream_names[idx]] = Position(timestamp=ts, skip=count + 1)
+                else:
+                    positions[stream_names[idx]] = Position(timestamp=event["timestamp"], skip=1)
+
+        if state == LogState.COMPLETE:
+            return state, last_description, last_describe_job_call
+
+        if state == LogState.JOB_COMPLETE:
+            state = LogState.COMPLETE
+        elif time.time() - last_describe_job_call >= 30:
+            description = await self.describe_training_job_async(job_name)
+            last_describe_job_call = time.time()
+
+            if await sync_to_async(secondary_training_status_changed)(description, last_description):
+                self.log.info(
+                    await sync_to_async(secondary_training_status_message)(description, last_description)
+                )
+                last_description = description
+
+            status = description["TrainingJobStatus"]
+
+            if status not in self.non_terminal_states:
+                state = LogState.JOB_COMPLETE
+        return state, last_description, last_describe_job_call
+
+    async def get_multi_stream(
+        self, log_group: str, streams: list[str], positions: dict[str, Any]
+    ) -> AsyncGenerator[Any, tuple[int, Any | None]]:
+        """Iterate over the available events coming and interleaving the events from each stream so they're yielded in timestamp order.
+
+        :param log_group: The name of the log group.
+        :param streams: A list of the log stream names. The position of the stream in this list is
+            the stream number.
+        :param positions: A list of pairs of (timestamp, skip) which represents the last record
+            read from each stream.
+        """
+        positions = positions or {s: Position(timestamp=0, skip=0) for s in streams}
+        events: list[Any | None] = []
+
+        logs_hook = AwsLogsHook(aws_conn_id=self.aws_conn_id, region_name=self.region_name)
+        event_iters = [
+            logs_hook.get_log_events_async(log_group, s, positions[s].timestamp, positions[s].skip)
+            for s in streams
+        ]
+        for event_stream in event_iters:
+            if not event_stream:
+                events.append(None)
+                continue
+
+            try:
+                events.append(await event_stream.__anext__())
+            except StopAsyncIteration:
+                events.append(None)
+
+            while any(events):
+                i = argmin(events, lambda x: x["timestamp"] if x else 9999999999) or 0
+                yield i, events[i]
+
+                try:
+                    events[i] = await event_iters[i].__anext__()
+                except StopAsyncIteration:
+                    events[i] = None

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

@@ -25,6 +25,7 @@ from functools import cached_property
 from typing import TYPE_CHECKING, Any, List, Sequence, cast
 
 from botocore.exceptions import ClientError, WaiterError
+from deprecated import deprecated
 
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
@@ -263,13 +264,14 @@ class EksCreateClusterOperator(BaseOperator):
         return EksHook(aws_conn_id=self.aws_conn_id, region_name=self.region)
 
     @property
-    def eks_hook(self):
-        warnings.warn(
+    @deprecated(
+        reason=(
             "`eks_hook` property is deprecated and will be removed in the future. "
-            "Please use `hook` property instead.",
-            AirflowProviderDeprecationWarning,
-            stacklevel=2,
-        )
+            "Please use `hook` property instead."
+        ),
+        category=AirflowProviderDeprecationWarning,
+    )
+    def eks_hook(self):
         return self.hook
 
     def execute(self, context: Context):

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

@@ -0,0 +1,218 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from airflow.configuration import conf
+from airflow.providers.amazon.aws.hooks.neptune import NeptuneHook
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune import (
+    NeptuneClusterAvailableTrigger,
+    NeptuneClusterStoppedTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
+
+class NeptuneStartDbClusterOperator(AwsBaseOperator[NeptuneHook]):
+    """Starts an Amazon Neptune DB cluster.
+
+    Amazon Neptune Database is a serverless graph database designed for superior scalability
+    and availability. Neptune Database provides built-in security, continuous backups, and
+    integrations with other AWS services
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:NeptuneStartDbClusterOperator`
+
+    :param db_cluster_id: The DB cluster identifier of the Neptune DB cluster to be started.
+    :param wait_for_completion: Whether to wait for the cluster to start. (default: True)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to start.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param waiter_delay: Time in seconds to wait between status checks.
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    :return: dictionary with Neptune cluster id
+    """
+
+    aws_hook_class = NeptuneHook
+    template_fields: Sequence[str] = aws_template_fields("cluster_id")
+
+    def __init__(
+        self,
+        db_cluster_id: str,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_id = db_cluster_id
+        self.wait_for_completion = wait_for_completion
+        self.deferrable = deferrable
+        self.delay = waiter_delay
+        self.max_attempts = waiter_max_attempts
+
+    def execute(self, context: Context) -> dict[str, str]:
+        self.log.info("Starting Neptune cluster: %s", self.cluster_id)
+
+        # Check to make sure the cluster is not already available.
+        status = self.hook.get_cluster_status(self.cluster_id)
+        if status.lower() in NeptuneHook.AVAILABLE_STATES:
+            self.log.info("Neptune cluster %s is already available.", self.cluster_id)
+            return {"db_cluster_id": self.cluster_id}
+
+        resp = self.hook.conn.start_db_cluster(DBClusterIdentifier=self.cluster_id)
+        status = resp.get("DBClusters", {}).get("Status", "Unknown")
+
+        if self.deferrable:
+            self.log.info("Deferring for cluster start: %s", self.cluster_id)
+
+            self.defer(
+                trigger=NeptuneClusterAvailableTrigger(
+                    aws_conn_id=self.aws_conn_id,
+                    db_cluster_id=self.cluster_id,
+                    waiter_delay=self.delay,
+                    waiter_max_attempts=self.max_attempts,
+                ),
+                method_name="execute_complete",
+            )
+
+        elif self.wait_for_completion:
+            self.log.info("Waiting for Neptune cluster %s to start.", self.cluster_id)
+            self.hook.wait_for_cluster_availability(self.cluster_id, self.delay, self.max_attempts)
+
+        return {"db_cluster_id": self.cluster_id}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, str]:
+        status = ""
+        cluster_id = ""
+
+        if event:
+            status = event.get("status", "")
+            cluster_id = event.get("cluster_id", "")
+
+        self.log.info("Neptune cluster %s available with status: %s", cluster_id, status)
+
+        return {"db_cluster_id": cluster_id}
+
+
+class NeptuneStopDbClusterOperator(AwsBaseOperator[NeptuneHook]):
+    """
+    Stops an Amazon Neptune DB cluster.
+
+    Amazon Neptune Database is a serverless graph database designed for superior scalability
+    and availability. Neptune Database provides built-in security, continuous backups, and
+    integrations with other AWS services
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:NeptuneStartDbClusterOperator`
+
+    :param db_cluster_id: The DB cluster identifier of the Neptune DB cluster to be stopped.
+    :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
+    :param deferrable: If True, the operator will wait asynchronously for the cluster to stop.
+        This implies waiting for completion. This mode requires aiobotocore module to be installed.
+        (default: False)
+    :param waiter_delay: Time in seconds to wait between status checks.
+    :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+    :param aws_conn_id: The Airflow connection used for AWS credentials.
+        If this is ``None`` or empty then the default boto3 behaviour is used. If
+        running Airflow in a distributed manner and aws_conn_id is None or
+        empty, then default boto3 configuration would be used (and must be
+        maintained on each worker node).
+    :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+    :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+        https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+    :return: dictionary with Neptune cluster id
+    """
+
+    aws_hook_class = NeptuneHook
+    template_fields: Sequence[str] = aws_template_fields("cluster_id")
+
+    def __init__(
+        self,
+        db_cluster_id: str,
+        wait_for_completion: bool = True,
+        waiter_delay: int = 30,
+        waiter_max_attempts: int = 60,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.cluster_id = db_cluster_id
+        self.wait_for_completion = wait_for_completion
+        self.deferrable = deferrable
+        self.delay = waiter_delay
+        self.max_attempts = waiter_max_attempts
+
+    def execute(self, context: Context) -> dict[str, str]:
+        self.log.info("Stopping Neptune cluster: %s", self.cluster_id)
+
+        # Check to make sure the cluster is not already stopped.
+        status = self.hook.get_cluster_status(self.cluster_id)
+        if status.lower() in NeptuneHook.STOPPED_STATES:
+            self.log.info("Neptune cluster %s is already stopped.", self.cluster_id)
+            return {"db_cluster_id": self.cluster_id}
+
+        resp = self.hook.conn.stop_db_cluster(DBClusterIdentifier=self.cluster_id)
+        status = resp.get("DBClusters", {}).get("Status", "Unknown")
+
+        if self.deferrable:
+            self.log.info("Deferring for cluster stop: %s", self.cluster_id)
+
+            self.defer(
+                trigger=NeptuneClusterStoppedTrigger(
+                    aws_conn_id=self.aws_conn_id,
+                    db_cluster_id=self.cluster_id,
+                    waiter_delay=self.delay,
+                    waiter_max_attempts=self.max_attempts,
+                ),
+                method_name="execute_complete",
+            )
+
+        elif self.wait_for_completion:
+            self.log.info("Waiting for Neptune cluster %s to start.", self.cluster_id)
+            self.hook.wait_for_cluster_stopped(self.cluster_id, self.delay, self.max_attempts)
+
+        return {"db_cluster_id": self.cluster_id}
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, str]:
+        status = ""
+        cluster_id = ""
+
+        if event:
+            status = event.get("status", "")
+            cluster_id = event.get("cluster_id", "")
+
+        self.log.info("Neptune cluster %s stopped with status: %s", cluster_id, status)
+
+        return {"db_cluster_id": cluster_id}

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

@@ -29,9 +29,14 @@ from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.models import BaseOperator
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
-from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
+from airflow.providers.amazon.aws.hooks.sagemaker import (
+    LogState,
+    SageMakerHook,
+    secondary_training_status_message,
+)
 from airflow.providers.amazon.aws.triggers.sagemaker import (
     SageMakerPipelineTrigger,
+    SageMakerTrainingPrintLogTrigger,
     SageMakerTrigger,
 )
 from airflow.providers.amazon.aws.utils import trim_none_values
@@ -899,9 +904,11 @@ class SageMakerTuningOperator(SageMakerBaseOperator):
                     aws_conn_id=self.aws_conn_id,
                 ),
                 method_name="execute_complete",
-                timeout=datetime.timedelta(seconds=self.max_ingestion_time)
-                if self.max_ingestion_time is not None
-                else None,
+                timeout=(
+                    datetime.timedelta(seconds=self.max_ingestion_time)
+                    if self.max_ingestion_time is not None
+                    else None
+                ),
             )
             description = {}  # never executed but makes static checkers happy
         elif self.wait_for_completion:
@@ -1085,28 +1092,80 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
             raise AirflowException(f"Sagemaker Training Job creation failed: {response}")
 
         if self.deferrable and self.wait_for_completion:
-            self.defer(
-                timeout=self.execution_timeout,
-                trigger=SageMakerTrigger(
+            description = self.hook.describe_training_job(self.config["TrainingJobName"])
+            status = description["TrainingJobStatus"]
+
+            if self.print_log:
+                instance_count = description["ResourceConfig"]["InstanceCount"]
+                last_describe_job_call = time.monotonic()
+                job_already_completed = status not in self.hook.non_terminal_states
+                _, description, last_describe_job_call = self.hook.describe_training_job_with_log(
+                    self.config["TrainingJobName"],
+                    {},
+                    [],
+                    instance_count,
+                    LogState.COMPLETE if job_already_completed else LogState.TAILING,
+                    description,
+                    last_describe_job_call,
+                )
+                self.log.info(secondary_training_status_message(description, None))
+
+            if status in self.hook.failed_states:
+                reason = description.get("FailureReason", "(No reason provided)")
+                raise AirflowException(f"SageMaker job failed because {reason}")
+            elif status == "Completed":
+                log_message = f"{self.task_id} completed successfully."
+                if self.print_log:
+                    billable_seconds = SageMakerHook.count_billable_seconds(
+                        training_start_time=description["TrainingStartTime"],
+                        training_end_time=description["TrainingEndTime"],
+                        instance_count=instance_count,
+                    )
+                    log_message = f"Billable seconds: {billable_seconds}\n{log_message}"
+                self.log.info(log_message)
+                return {"Training": serialize(description)}
+
+            timeout = self.execution_timeout
+            if self.max_ingestion_time:
+                timeout = datetime.timedelta(seconds=self.max_ingestion_time)
+
+            trigger: SageMakerTrainingPrintLogTrigger | SageMakerTrigger
+            if self.print_log:
+                trigger = SageMakerTrainingPrintLogTrigger(
+                    job_name=self.config["TrainingJobName"],
+                    poke_interval=self.check_interval,
+                    aws_conn_id=self.aws_conn_id,
+                )
+            else:
+                trigger = SageMakerTrigger(
                     job_name=self.config["TrainingJobName"],
                     job_type="Training",
                     poke_interval=self.check_interval,
                     max_attempts=self.max_attempts,
                     aws_conn_id=self.aws_conn_id,
-                ),
+                )
+
+            self.defer(
+                timeout=timeout,
+                trigger=trigger,
                 method_name="execute_complete",
             )
 
-        self.serialized_training_data = serialize(
-            self.hook.describe_training_job(self.config["TrainingJobName"])
-        )
-        return {"Training": self.serialized_training_data}
+        return self.serialize_result()
+
+    def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, dict]:
+        if event is None:
+            err_msg = "Trigger error: event is None"
+            self.log.error(err_msg)
+            raise AirflowException(err_msg)
 
-    def execute_complete(self, context, event=None):
         if event["status"] != "success":
             raise AirflowException(f"Error while running job: {event}")
-        else:
-            self.log.info(event["message"])
+
+        self.log.info(event["message"])
+        return self.serialize_result()
+
+    def serialize_result(self) -> dict[str, dict]:
         self.serialized_training_data = serialize(
             self.hook.describe_training_job(self.config["TrainingJobName"])
         )

airflow/providers/amazon/aws/sensors/batch.py

@@ -60,8 +60,8 @@ class BatchSensor(BaseSensorOperator):
         aws_conn_id: str = "aws_default",
         region_name: str | None = None,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
-        poke_interval: float = 5,
-        max_retries: int = 5,
+        poke_interval: float = 30,
+        max_retries: int = 4200,
         **kwargs,
     ):
         super().__init__(**kwargs)