Add deferrable param in SageMakerTrainingOperator

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

@@ -28,6 +28,7 @@
 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.triggers.sagemaker import SageMakerTrigger
 from airflow.providers.amazon.aws.utils import trim_none_values
 from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
 from airflow.providers.amazon.aws.utils.tags import format_tags
@@ -683,6 +684,8 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
     :param print_log: if the operator should print the cloudwatch log during training
     :param check_interval: if wait is set to be true, this is the time interval
         in seconds which the operator will check the status of the training job
+    :param max_attempts: Number of times to poll for query state before returning the current state,
+        defaults to None.
     :param max_ingestion_time: If wait is set to True, the operation fails if the training job
         doesn't finish within max_ingestion_time seconds. If you set this parameter to None,
         the operation does not timeout.
@@ -691,6 +694,8 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
     :param action_if_job_exists: Behaviour if the job name already exists. Possible options are "timestamp"
         (default), "increment" (deprecated) and "fail".
         This is only relevant if check_if_job_exists is True.
+    :param deferrable: Run operator in the deferrable mode. This is only effective if wait_for_completion is
+        set to True.
     :return Dict: Returns The ARN of the training job created in Amazon SageMaker.
     """
 
@@ -702,15 +707,18 @@ def __init__(
         wait_for_completion: bool = True,
         print_log: bool = True,
         check_interval: int = CHECK_INTERVAL_SECOND,
+        max_attempts: int | None = None,
         max_ingestion_time: int | None = None,
         check_if_job_exists: bool = True,
         action_if_job_exists: str = "timestamp",
+        deferrable: bool = False,
         **kwargs,
     ):
         super().__init__(config=config, aws_conn_id=aws_conn_id, **kwargs)
         self.wait_for_completion = wait_for_completion
         self.print_log = print_log
         self.check_interval = check_interval
+        self.max_attempts = max_attempts or 60
         self.max_ingestion_time = max_ingestion_time
         self.check_if_job_exists = check_if_job_exists
         if action_if_job_exists in {"timestamp", "increment", "fail"}:
@@ -727,6 +735,7 @@ def __init__(
                 f"Argument action_if_job_exists accepts only 'timestamp', 'increment' and 'fail'. \
                 Provided value: '{action_if_job_exists}'."
             )
+        self.deferrable = deferrable
 
     def expand_role(self) -> None:
         """Expands an IAM role name into an ARN."""
@@ -753,17 +762,50 @@ def execute(self, context: Context) -> dict:
             )
 
         self.log.info("Creating SageMaker training job %s.", self.config["TrainingJobName"])
+
+        if self.deferrable and not self.wait_for_completion:
+            self.log.warning(
+                "Setting deferrable to True does not have effect when wait_for_completion is set to False."
+            )
+
+        wait_for_completion = self.wait_for_completion
+        if self.deferrable and self.wait_for_completion:
+            # Set wait_for_completion to False so that it waits for the status in the deferred task.
+            wait_for_completion = False
+
         response = self.hook.create_training_job(
             self.config,
-            wait_for_completion=self.wait_for_completion,
+            wait_for_completion=wait_for_completion,
             print_log=self.print_log,
             check_interval=self.check_interval,
             max_ingestion_time=self.max_ingestion_time,
         )
         if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
             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(
+                    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,
+                ),
+                method_name="execute_complete",
+            )
+
+        result = {"Training": serialize(self.hook.describe_training_job(self.config["TrainingJobName"]))}
+        return result
+
+    def execute_complete(self, context, event=None):
+        if event["status"] != "success":
+            raise AirflowException(f"Error while running job: {event}")
         else:
-            return {"Training": serialize(self.hook.describe_training_job(self.config["TrainingJobName"]))}
+            self.log.info(event["message"])
+        result = {"Training": serialize(self.hook.describe_training_job(self.config["TrainingJobName"]))}
+        return result
 
 
 class SageMakerDeleteModelOperator(SageMakerBaseOperator):

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

@@ -0,0 +1,101 @@
+# 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 Any
+
+from airflow.compat.functools import cached_property
+from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+
+class SageMakerTrigger(BaseTrigger):
+    """
+    SageMakerTrigger is fired as deferred class with params to run the task in triggerer.
+
+    :param job_name: name of the job to check status
+    :param job_type: Type of the sagemaker job whether it is Transform or Training
+    :param poke_interval:  polling period in seconds to check for the status
+    :param max_attempts: Number of times to poll for query state before returning the current state,
+        defaults to None.
+    :param aws_conn_id: AWS connection ID for sagemaker
+    """
+
+    def __init__(
+        self,
+        job_name: str,
+        job_type: str,
+        poke_interval: int = 30,
+        max_attempts: int | None = None,
+        aws_conn_id: str = "aws_default",
+    ):
+        super().__init__()
+        self.job_name = job_name
+        self.job_type = job_type
+        self.poke_interval = poke_interval
+        self.max_attempts = max_attempts
+        self.aws_conn_id = aws_conn_id
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serializes SagemakerTrigger arguments and classpath."""
+        return (
+            "airflow.providers.amazon.aws.triggers.sagemaker.SageMakerTrigger",
+            {
+                "job_name": self.job_name,
+                "job_type": self.job_type,
+                "poke_interval": self.poke_interval,
+                "max_attempts": self.max_attempts,
+                "aws_conn_id": self.aws_conn_id,
+            },
+        )
+
+    @cached_property
+    def hook(self) -> SageMakerHook:
+        return SageMakerHook(aws_conn_id=self.aws_conn_id)
+
+    @staticmethod
+    def _get_job_type_waiter(job_type: str) -> str:
+        return {
+            "training": "TrainingJobComplete",
+            "transform": "TransformJobComplete",
+            "processing": "ProcessingJobComplete",
+        }[job_type.lower()]
+
+    @staticmethod
+    def _get_job_type_waiter_job_name_arg(job_type: str) -> str:
+        return {
+            "training": "TrainingJobName",
+            "transform": "TransformJobName",
+            "processing": "ProcessingJobName",
+        }[job_type.lower()]
+
+    async def run(self):
+        self.log.info("job name is %s and job type is %s", self.job_name, self.job_type)
+        async with self.hook.async_conn as client:
+            waiter = self.hook.get_waiter(
+                self._get_job_type_waiter(self.job_type), deferrable=True, client=client
+            )
+            waiter_args = {
+                self._get_job_type_waiter_job_name_arg(self.job_type): self.job_name,
+                "WaiterConfig": {
+                    "Delay": self.poke_interval,
+                    "MaxAttempts": self.max_attempts,
+                },
+            }
+            await waiter.wait(**waiter_args)
+        yield TriggerEvent({"status": "success", "message": "Job completed."})

airflow/providers/amazon/aws/waiters/sagemaker.json

@@ -0,0 +1,83 @@
+{
+    "version": 2,
+    "waiters": {
+        "TrainingJobComplete": {
+            "delay": 30,
+            "operation": "DescribeTrainingJob",
+            "maxAttempts": 60,
+            "description": "Wait until job is COMPLETED",
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "TrainingJobStatus",
+                    "expected": "Completed",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "TrainingJobStatus",
+                    "expected": "Failed",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "TrainingJobStatus",
+                    "expected": "Stopped",
+                    "state": "failure"
+                }
+            ]
+        },
+        "TransformJobComplete": {
+            "delay": 30,
+            "operation": "DescribeTransformJob",
+            "maxAttempts": 60,
+            "description": "Wait until job is COMPLETED",
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "TransformJobStatus",
+                    "expected": "Completed",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "TransformJobStatus",
+                    "expected": "Failed",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "TransformJobStatus",
+                    "expected": "Stopped",
+                    "state": "failure"
+                }
+            ]
+        },
+        "ProcessingJobComplete": {
+            "delay": 30,
+            "operation": "DescribeProcessingJob",
+            "maxAttempts": 60,
+            "description": "Wait until job is COMPLETED",
+            "acceptors": [
+                {
+                    "matcher": "path",
+                    "argument": "ProcessingJobStatus",
+                    "expected": "Completed",
+                    "state": "success"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "ProcessingJobStatus",
+                    "expected": "Failed",
+                    "state": "failure"
+                },
+                {
+                    "matcher": "path",
+                    "argument": "ProcessingJobStatus",
+                    "expected": "Stopped",
+                    "state": "failure"
+                }
+            ]
+        }
+    }
+}

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

@@ -21,10 +21,11 @@
 import pytest
 from botocore.exceptions import ClientError
 
-from airflow.exceptions import AirflowException
+from airflow.exceptions import AirflowException, TaskDeferred
 from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
 from airflow.providers.amazon.aws.operators import sagemaker
 from airflow.providers.amazon.aws.operators.sagemaker import SageMakerTrainingOperator
+from airflow.providers.amazon.aws.triggers.sagemaker import SageMakerTrigger
 
 EXPECTED_INTEGER_FIELDS: list[list[str]] = [
     ["ResourceConfig", "InstanceCount"],
@@ -113,3 +114,16 @@ def test_execute_with_failure(self, mock_training, mock_desc):
         }
         with pytest.raises(AirflowException):
             self.sagemaker.execute(None)
+
+    @mock.patch.object(SageMakerHook, "create_training_job")
+    def test_operator_defer(self, mock_training):
+        mock_training.return_value = {
+            "TrainingJobArn": "test_arn",
+            "ResponseMetadata": {"HTTPStatusCode": 200},
+        }
+        self.sagemaker.deferrable = True
+        self.sagemaker.wait_for_completion = True
+        self.sagemaker.check_if_job_exists = False
+        with pytest.raises(TaskDeferred) as exc:
+            self.sagemaker.execute(context=None)
+        assert isinstance(exc.value.trigger, SageMakerTrigger), "Trigger is not a SagemakerTrigger"