diff --git a/airflow/providers/amazon/aws/hooks/sagemaker.py b/airflow/providers/amazon/aws/hooks/sagemaker.py index ae57e832d1b90..6653e0eb6d7f6 100644 --- a/airflow/providers/amazon/aws/hooks/sagemaker.py +++ b/airflow/providers/amazon/aws/hooks/sagemaker.py @@ -24,6 +24,7 @@ import tempfile import time import warnings +from collections import Counter from datetime import datetime from functools import partial from typing import Any, Callable, Generator, cast @@ -146,6 +147,7 @@ class SageMakerHook(AwsBaseHook): non_terminal_states = {"InProgress", "Stopping"} endpoint_non_terminal_states = {"Creating", "Updating", "SystemUpdating", "RollingBack", "Deleting"} + pipeline_non_terminal_states = {"Executing", "Stopping"} failed_states = {"Failed"} def __init__(self, *args, **kwargs): @@ -654,22 +656,21 @@ def check_status( check_interval: int, max_ingestion_time: int | None = None, non_terminal_states: set | None = None, - ): + ) -> dict: """ - Check status of a SageMaker job + Check status of a SageMaker resource - :param job_name: name of the job to check status - :param key: the key of the response dict - that points to the state + :param job_name: name of the resource to check status, can be a job but also pipeline for instance. + :param key: the key of the response dict that points to the state :param describe_function: the function used to retrieve the status :param args: the arguments for the function :param check_interval: the time interval in seconds which the operator - will check the status of any SageMaker job + will check the status of any SageMaker resource :param max_ingestion_time: the maximum ingestion time in seconds. Any - SageMaker jobs that run longer than this will fail. Setting this to - None implies no timeout for any SageMaker job. + SageMaker resources that run longer than this will fail. Setting this to + None implies no timeout for any SageMaker resource. :param non_terminal_states: the set of nonterminal states - :return: response of describe call after job is done + :return: response of describe call after resource is done """ if not non_terminal_states: non_terminal_states = self.non_terminal_states @@ -683,22 +684,22 @@ def check_status( try: response = describe_function(job_name) status = response[key] - self.log.info("Job still running for %s seconds... current status is %s", sec, status) + self.log.info("Resource still running for %s seconds... current status is %s", sec, status) except KeyError: - raise AirflowException("Could not get status of the SageMaker job") + raise AirflowException("Could not get status of the SageMaker resource") except ClientError: raise AirflowException("AWS request failed, check logs for more info") if status in self.failed_states: - raise AirflowException(f"SageMaker job failed because {response['FailureReason']}") + raise AirflowException(f"SageMaker resource failed because {response['FailureReason']}") elif status not in non_terminal_states: break if max_ingestion_time and sec > max_ingestion_time: - # ensure that the job gets killed if the max ingestion time is exceeded - raise AirflowException(f"SageMaker job took more than {max_ingestion_time} seconds") + # ensure that the resource gets killed if the max ingestion time is exceeded + raise AirflowException(f"SageMaker resource took more than {max_ingestion_time} seconds") - self.log.info("SageMaker Job completed") + self.log.info("SageMaker resource completed") return response def check_training_status_with_log( @@ -1010,3 +1011,122 @@ def delete_model(self, model_name: str): except Exception as general_error: self.log.error("Failed to delete model, error: %s", general_error) raise + + def describe_pipeline_exec(self, pipeline_exec_arn: str, verbose: bool = False): + """Get info about a SageMaker pipeline execution + + :param pipeline_exec_arn: arn of the pipeline execution + :param verbose: Whether to log details about the steps status in the pipeline execution + """ + if verbose: + res = self.conn.list_pipeline_execution_steps(PipelineExecutionArn=pipeline_exec_arn) + count_by_state = Counter(s["StepStatus"] for s in res["PipelineExecutionSteps"]) + running_steps = [ + s["StepName"] for s in res["PipelineExecutionSteps"] if s["StepStatus"] == "Executing" + ] + self.log.info("state of the pipeline steps: %s", count_by_state) + self.log.info("steps currently in progress: %s", running_steps) + + return self.conn.describe_pipeline_execution(PipelineExecutionArn=pipeline_exec_arn) + + def start_pipeline( + self, + pipeline_name: str, + display_name: str = "airflow-triggered-execution", + pipeline_params: dict | None = None, + wait_for_completion: bool = False, + check_interval: int = 30, + verbose: bool = True, + ) -> str: + """ + Start a new execution for a SageMaker pipeline + + :param pipeline_name: Name of the pipeline to start (this is _not_ the ARN). + :param display_name: The name this pipeline execution will have in the UI. Doesn't need to be unique. + :param pipeline_params: Optional parameters for the pipeline. + All parameters supplied need to already be present in the pipeline definition. + :param wait_for_completion: Will only return once the pipeline is complete if true. + :param check_interval: How long to wait between checks for pipeline status when waiting for + completion. + :param verbose: Whether to print steps details when waiting for completion. + Defaults to true, consider turning off for pipelines that have thousands of steps. + + :return: the ARN of the pipeline execution launched. + """ + if pipeline_params is None: + pipeline_params = {} + formatted_params = [{"Name": kvp[0], "Value": kvp[1]} for kvp in pipeline_params.items()] + + try: + res = self.conn.start_pipeline_execution( + PipelineName=pipeline_name, + PipelineExecutionDisplayName=display_name, + PipelineParameters=formatted_params, + ) + except ClientError as ce: + self.log.error("Failed to start pipeline execution, error: %s", ce) + raise + + arn = res["PipelineExecutionArn"] + if wait_for_completion: + self.check_status( + arn, + "PipelineExecutionStatus", + lambda p: self.describe_pipeline_exec(p, verbose), + check_interval, + non_terminal_states=self.pipeline_non_terminal_states, + ) + return arn + + def stop_pipeline( + self, + pipeline_exec_arn: str, + wait_for_completion: bool = False, + check_interval: int = 10, + verbose: bool = True, + fail_if_not_running: bool = False, + ) -> str: + """Stop SageMaker pipeline execution + + :param pipeline_exec_arn: Amazon Resource Name (ARN) of the pipeline execution. + It's the ARN of the pipeline itself followed by "/execution/" and an id. + :param wait_for_completion: Whether to wait for the pipeline to reach a final state. + (i.e. either 'Stopped' or 'Failed') + :param check_interval: How long to wait between checks for pipeline status when waiting for + completion. + :param verbose: Whether to print steps details when waiting for completion. + Defaults to true, consider turning off for pipelines that have thousands of steps. + :param fail_if_not_running: This method will raise an exception if the pipeline we're trying to stop + is not in an "Executing" state when the call is sent (which would mean that the pipeline is + already either stopping or stopped). + Note that setting this to True will raise an error if the pipeline finished successfully before it + was stopped. + :return: Status of the pipeline execution after the operation. + One of 'Executing'|'Stopping'|'Stopped'|'Failed'|'Succeeded'. + """ + try: + self.conn.stop_pipeline_execution(PipelineExecutionArn=pipeline_exec_arn) + except ClientError as ce: + # we have to rely on the message to catch the right error here, because its type + # (ValidationException) is shared with other kinds of error (for instance, badly formatted ARN) + if ( + not fail_if_not_running + and "Only pipelines with 'Executing' status can be stopped" in ce.response["Error"]["Message"] + ): + self.log.warning("Cannot stop pipeline execution, as it was not running: %s", ce) + else: + self.log.error(ce) + raise + + res = self.describe_pipeline_exec(pipeline_exec_arn) + + if wait_for_completion and res["PipelineExecutionStatus"] in self.pipeline_non_terminal_states: + res = self.check_status( + pipeline_exec_arn, + "PipelineExecutionStatus", + lambda p: self.describe_pipeline_exec(p, verbose), + check_interval, + non_terminal_states=self.pipeline_non_terminal_states, + ) + + return res["PipelineExecutionStatus"] diff --git a/airflow/providers/amazon/aws/operators/sagemaker.py b/airflow/providers/amazon/aws/operators/sagemaker.py index c5f08db049206..4b969002b339e 100644 --- a/airflow/providers/amazon/aws/operators/sagemaker.py +++ b/airflow/providers/amazon/aws/operators/sagemaker.py @@ -103,7 +103,7 @@ def _create_integer_fields(self) -> None: """ self.integer_fields = [] - def execute(self, context: Context) -> None | dict: + def execute(self, context: Context): raise NotImplementedError("Please implement execute() in sub class!") @cached_property @@ -750,3 +750,121 @@ def execute(self, context: Context) -> Any: sagemaker_hook = SageMakerHook(aws_conn_id=self.aws_conn_id) sagemaker_hook.delete_model(model_name=self.config["ModelName"]) self.log.info("Model %s deleted successfully.", self.config["ModelName"]) + + +class SageMakerStartPipelineOperator(SageMakerBaseOperator): + """ + Starts a SageMaker pipeline execution. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:SageMakerStartPipelineOperator` + + :param config: The configuration to start the pipeline execution. + :param aws_conn_id: The AWS connection ID to use. + :param pipeline_name: Name of the pipeline to start. + :param display_name: The name this pipeline execution will have in the UI. Doesn't need to be unique. + :param pipeline_params: Optional parameters for the pipeline. + All parameters supplied need to already be present in the pipeline definition. + :param wait_for_completion: If true, this operator will only complete once the pipeline is complete. + :param check_interval: How long to wait between checks for pipeline status when waiting for completion. + :param verbose: Whether to print steps details when waiting for completion. + Defaults to true, consider turning off for pipelines that have thousands of steps. + + :return str: Returns The ARN of the pipeline execution created in Amazon SageMaker. + """ + + template_fields: Sequence[str] = ("aws_conn_id", "pipeline_name", "display_name", "pipeline_params") + + def __init__( + self, + *, + aws_conn_id: str = DEFAULT_CONN_ID, + pipeline_name: str, + display_name: str = "airflow-triggered-execution", + pipeline_params: dict | None = None, + wait_for_completion: bool = False, + check_interval: int = CHECK_INTERVAL_SECOND, + verbose: bool = True, + **kwargs, + ): + super().__init__(config={}, aws_conn_id=aws_conn_id, **kwargs) + self.pipeline_name = pipeline_name + self.display_name = display_name + self.pipeline_params = pipeline_params + self.wait_for_completion = wait_for_completion + self.check_interval = check_interval + self.verbose = verbose + + def execute(self, context: Context) -> str: + arn = self.hook.start_pipeline( + pipeline_name=self.pipeline_name, + display_name=self.display_name, + pipeline_params=self.pipeline_params, + wait_for_completion=self.wait_for_completion, + check_interval=self.check_interval, + verbose=self.verbose, + ) + self.log.info( + "Starting a new execution for pipeline %s, running with ARN %s", self.pipeline_name, arn + ) + return arn + + +class SageMakerStopPipelineOperator(SageMakerBaseOperator): + """ + Stops a SageMaker pipeline execution. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:SageMakerStopPipelineOperator` + + :param config: The configuration to start the pipeline execution. + :param aws_conn_id: The AWS connection ID to use. + :param pipeline_exec_arn: Amazon Resource Name of the pipeline execution to stop. + :param wait_for_completion: If true, this operator will only complete once the pipeline is fully stopped. + :param check_interval: How long to wait between checks for pipeline status when waiting for completion. + :param verbose: Whether to print steps details when waiting for completion. + Defaults to true, consider turning off for pipelines that have thousands of steps. + :param fail_if_not_running: raises an exception if the pipeline stopped or succeeded before this was run + + :return str: Returns the status of the pipeline execution after the operation has been done. + """ + + template_fields: Sequence[str] = ( + "aws_conn_id", + "pipeline_exec_arn", + ) + + def __init__( + self, + *, + aws_conn_id: str = DEFAULT_CONN_ID, + pipeline_exec_arn: str, + wait_for_completion: bool = False, + check_interval: int = CHECK_INTERVAL_SECOND, + verbose: bool = True, + fail_if_not_running: bool = False, + **kwargs, + ): + super().__init__(config={}, aws_conn_id=aws_conn_id, **kwargs) + self.pipeline_exec_arn = pipeline_exec_arn + self.wait_for_completion = wait_for_completion + self.check_interval = check_interval + self.verbose = verbose + self.fail_if_not_running = fail_if_not_running + + def execute(self, context: Context) -> str: + status = self.hook.stop_pipeline( + pipeline_exec_arn=self.pipeline_exec_arn, + wait_for_completion=self.wait_for_completion, + check_interval=self.check_interval, + verbose=self.verbose, + fail_if_not_running=self.fail_if_not_running, + ) + self.log.info( + "Stop requested for pipeline execution with ARN %s. Status is now %s", + self.pipeline_exec_arn, + status, + ) + return status diff --git a/airflow/providers/amazon/aws/sensors/sagemaker.py b/airflow/providers/amazon/aws/sensors/sagemaker.py index 2d9c9aacf0f06..a93135b48ce17 100644 --- a/airflow/providers/amazon/aws/sensors/sagemaker.py +++ b/airflow/providers/amazon/aws/sensors/sagemaker.py @@ -37,9 +37,10 @@ class SageMakerBaseSensor(BaseSensorOperator): ui_color = "#ededed" - def __init__(self, *, aws_conn_id: str = "aws_default", **kwargs): + def __init__(self, *, aws_conn_id: str = "aws_default", resource_type: str = "job", **kwargs): super().__init__(**kwargs) self.aws_conn_id = aws_conn_id + self.resource_type = resource_type # only used for logs, to say what kind of resource we are sensing self.hook: SageMakerHook | None = None def get_hook(self) -> SageMakerHook: @@ -55,12 +56,14 @@ def poke(self, context: Context): self.log.info("Bad HTTP response: %s", response) return False state = self.state_from_response(response) - self.log.info("Job currently %s", state) + self.log.info("%s currently %s", self.resource_type, state) if state in self.non_terminal_states(): return False if state in self.failed_states(): failed_reason = self.get_failed_reason_from_response(response) - raise AirflowException(f"Sagemaker job failed for the following reason: {failed_reason}") + raise AirflowException( + f"Sagemaker {self.resource_type} failed for the following reason: {failed_reason}" + ) return True def non_terminal_states(self) -> set[str]: @@ -269,3 +272,38 @@ def get_failed_reason_from_response(self, response): def state_from_response(self, response): return response["TrainingJobStatus"] + + +class SageMakerPipelineSensor(SageMakerBaseSensor): + """ + Polls the pipeline until it reaches a terminal state. Raises an + AirflowException with the failure reason if a failed state is reached. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto/sensor:SageMakerPipelineSensor` + + :param pipeline_exec_arn: ARN of the pipeline to watch. + :param verbose: Whether to print steps details while waiting for completion. + Defaults to true, consider turning off for pipelines that have thousands of steps. + """ + + template_fields: Sequence[str] = ("pipeline_exec_arn",) + + def __init__(self, *, pipeline_exec_arn: str, verbose: bool = True, **kwargs): + super().__init__(resource_type="pipeline", **kwargs) + self.pipeline_exec_arn = pipeline_exec_arn + self.verbose = verbose + + def non_terminal_states(self) -> set[str]: + return SageMakerHook.pipeline_non_terminal_states + + def failed_states(self) -> set[str]: + return SageMakerHook.failed_states + + def get_sagemaker_response(self) -> dict: + self.log.info("Poking Sagemaker Pipeline Execution %s", self.pipeline_exec_arn) + return self.get_hook().describe_pipeline_exec(self.pipeline_exec_arn, self.verbose) + + def state_from_response(self, response: dict) -> str: + return response["PipelineExecutionStatus"] diff --git a/docs/apache-airflow-providers-amazon/operators/sagemaker.rst b/docs/apache-airflow-providers-amazon/operators/sagemaker.rst index c929c954c8601..7023f2f5906a7 100644 --- a/docs/apache-airflow-providers-amazon/operators/sagemaker.rst +++ b/docs/apache-airflow-providers-amazon/operators/sagemaker.rst @@ -24,7 +24,7 @@ machine learning service. With Amazon SageMaker, data scientists and developers can quickly build and train machine learning models, and then deploy them into a production-ready hosted environment. -Airflow provides operators to create and interact with SageMaker Jobs. +Airflow provides operators to create and interact with SageMaker Jobs and Pipelines. Prerequisite Tasks ------------------ @@ -146,6 +146,34 @@ To create an Amazon Sagemaker endpoint you can use :start-after: [START howto_operator_sagemaker_endpoint] :end-before: [END howto_operator_sagemaker_endpoint] +.. _howto/operator:SageMakerStartPipelineOperator: + +Start an Amazon SageMaker pipeline execution +============================================ + +To trigger an execution run for an already-defined Amazon Sagemaker pipeline, you can use +:class:`~airflow.providers.amazon.aws.operators.sagemaker.SageMakerStartPipelineOperator`. + +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_sagemaker.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_sagemaker_start_pipeline] + :end-before: [END howto_operator_sagemaker_start_pipeline] + +.. _howto/operator:SageMakerStopPipelineOperator: + +Stop an Amazon SageMaker pipeline execution +=========================================== + +To stop an Amazon Sagemaker pipeline execution that is currently running, you can use +:class:`~airflow.providers.amazon.aws.operators.sagemaker.SageMakerStopPipelineOperator`. + +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_sagemaker.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_sagemaker_stop_pipeline] + :end-before: [END howto_operator_sagemaker_stop_pipeline] + Sensors ------- @@ -205,6 +233,20 @@ you can use :class:`~airflow.providers.amazon.aws.sensors.sagemaker.SageMakerEnd :start-after: [START howto_sensor_sagemaker_endpoint] :end-before: [END howto_sensor_sagemaker_endpoint] +.. _howto/sensor:SageMakerPipelineSensor: + +Wait on an Amazon SageMaker pipeline execution state +==================================================== + +To check the state of an Amazon Sagemaker pipeline execution until it reaches a terminal state +you can use :class:`~airflow.providers.amazon.aws.sensors.sagemaker.SageMakerPipelineSensor`. + +.. exampleinclude:: /../../tests/system/providers/amazon/aws/example_sagemaker.py + :language: python + :dedent: 4 + :start-after: [START howto_sensor_sagemaker_pipeline] + :end-before: [END howto_sensor_sagemaker_pipeline] + Reference --------- diff --git a/tests/providers/amazon/aws/hooks/test_sagemaker.py b/tests/providers/amazon/aws/hooks/test_sagemaker.py index 52906031d7a4f..fd12cf2796cd7 100644 --- a/tests/providers/amazon/aws/hooks/test_sagemaker.py +++ b/tests/providers/amazon/aws/hooks/test_sagemaker.py @@ -727,3 +727,98 @@ def test_delete_model_when_not_exist(self): ex = raised_exception.value assert ex.operation_name == "DeleteModel" assert ex.response["ResponseMetadata"]["HTTPStatusCode"] == 404 + + @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn", new_callable=mock.PropertyMock) + def test_start_pipeline_returns_arn(self, mock_conn): + mock_conn().start_pipeline_execution.return_value = {"PipelineExecutionArn": "hellotest"} + + hook = SageMakerHook(aws_conn_id="aws_default") + params_dict = {"one": "1", "two": "2"} + arn = hook.start_pipeline(pipeline_name="test_name", pipeline_params=params_dict) + + assert arn == "hellotest" + + args_passed = mock_conn().start_pipeline_execution.call_args[1] + assert args_passed["PipelineName"] == "test_name" + + # check conversion to the weird format for passing parameters (list of tuples) + assert len(args_passed["PipelineParameters"]) == 2 + for transformed_param in args_passed["PipelineParameters"]: + assert "Name" in transformed_param.keys() + assert "Value" in transformed_param.keys() + # Name contains the key + assert transformed_param["Name"] in params_dict.keys() + # Value contains the value associated with the key in Name + assert transformed_param["Value"] == params_dict[transformed_param["Name"]] + + @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn", new_callable=mock.PropertyMock) + def test_start_pipeline_waits_for_completion(self, mock_conn): + mock_conn().describe_pipeline_execution.side_effect = [ + {"PipelineExecutionStatus": "Executing"}, + {"PipelineExecutionStatus": "Executing"}, + {"PipelineExecutionStatus": "Succeeded"}, + ] + + hook = SageMakerHook(aws_conn_id="aws_default") + hook.start_pipeline(pipeline_name="test_name", wait_for_completion=True, check_interval=0) + + assert mock_conn().describe_pipeline_execution.call_count == 3 + + @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn", new_callable=mock.PropertyMock) + def test_stop_pipeline_returns_status(self, mock_conn): + mock_conn().describe_pipeline_execution.return_value = {"PipelineExecutionStatus": "Stopping"} + + hook = SageMakerHook(aws_conn_id="aws_default") + pipeline_status = hook.stop_pipeline(pipeline_exec_arn="test") + + assert pipeline_status == "Stopping" + mock_conn().stop_pipeline_execution.assert_called_once_with(PipelineExecutionArn="test") + + @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn", new_callable=mock.PropertyMock) + def test_stop_pipeline_waits_for_completion(self, mock_conn): + mock_conn().describe_pipeline_execution.side_effect = [ + {"PipelineExecutionStatus": "Stopping"}, + {"PipelineExecutionStatus": "Stopping"}, + {"PipelineExecutionStatus": "Stopped"}, + ] + + hook = SageMakerHook(aws_conn_id="aws_default") + pipeline_status = hook.stop_pipeline( + pipeline_exec_arn="test", wait_for_completion=True, check_interval=0 + ) + + assert pipeline_status == "Stopped" + assert mock_conn().describe_pipeline_execution.call_count == 3 + + @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn", new_callable=mock.PropertyMock) + def test_stop_pipeline_waits_for_completion_even_when_already_stopped(self, mock_conn): + mock_conn().stop_pipeline_execution.side_effect = ClientError( + error_response={"Error": {"Message": "Only pipelines with 'Executing' status can be stopped"}}, + operation_name="empty", + ) + mock_conn().describe_pipeline_execution.side_effect = [ + {"PipelineExecutionStatus": "Stopping"}, + {"PipelineExecutionStatus": "Stopped"}, + ] + + hook = SageMakerHook(aws_conn_id="aws_default") + pipeline_status = hook.stop_pipeline( + pipeline_exec_arn="test", wait_for_completion=True, check_interval=0 + ) + + assert pipeline_status == "Stopped" + + @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn", new_callable=mock.PropertyMock) + def test_stop_pipeline_raises_when_already_stopped_if_specified(self, mock_conn): + error = ClientError( + error_response={"Error": {"Message": "Only pipelines with 'Executing' status can be stopped"}}, + operation_name="empty", + ) + mock_conn().stop_pipeline_execution.side_effect = error + mock_conn().describe_pipeline_execution.return_value = {"PipelineExecutionStatus": "Stopping"} + + hook = SageMakerHook(aws_conn_id="aws_default") + with pytest.raises(ClientError) as raised_exception: + hook.stop_pipeline(pipeline_exec_arn="test", fail_if_not_running=True) + + assert raised_exception.value == error diff --git a/tests/providers/amazon/aws/operators/test_sagemaker_pipeline.py b/tests/providers/amazon/aws/operators/test_sagemaker_pipeline.py new file mode 100644 index 0000000000000..dde36274587e2 --- /dev/null +++ b/tests/providers/amazon/aws/operators/test_sagemaker_pipeline.py @@ -0,0 +1,69 @@ +# +# 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 unittest import mock + +from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook +from airflow.providers.amazon.aws.operators.sagemaker import ( + SageMakerStartPipelineOperator, + SageMakerStopPipelineOperator, +) + + +class TestSageMakerStartPipelineOperator: + @mock.patch.object(SageMakerHook, "start_pipeline") + def test_execute(self, start_pipeline): + op = SageMakerStartPipelineOperator( + task_id="test_sagemaker_operator", + pipeline_name="my_pipeline", + display_name="test_disp_name", + pipeline_params={"is_a_test": "yes"}, + wait_for_completion=True, + check_interval=12, + verbose=False, + ) + + op.execute(None) + + start_pipeline.assert_called_once_with( + pipeline_name="my_pipeline", + display_name="test_disp_name", + pipeline_params={"is_a_test": "yes"}, + wait_for_completion=True, + check_interval=12, + verbose=False, + ) + + +class TestSageMakerStopPipelineOperator: + @mock.patch.object(SageMakerHook, "stop_pipeline") + def test_execute(self, stop_pipeline): + op = SageMakerStopPipelineOperator( + task_id="test_sagemaker_operator", pipeline_exec_arn="pipeline_arn" + ) + + op.execute(None) + + stop_pipeline.assert_called_once_with( + pipeline_exec_arn="pipeline_arn", + wait_for_completion=False, + check_interval=30, + fail_if_not_running=False, + verbose=True, + ) diff --git a/tests/providers/amazon/aws/sensors/test_sagemaker_pipeline.py b/tests/providers/amazon/aws/sensors/test_sagemaker_pipeline.py new file mode 100644 index 0000000000000..ff8ddbf5530e9 --- /dev/null +++ b/tests/providers/amazon/aws/sensors/test_sagemaker_pipeline.py @@ -0,0 +1,86 @@ +# +# 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 + +import copy +from datetime import datetime +from unittest import mock + +import pytest + +from airflow.exceptions import AirflowException +from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook +from airflow.providers.amazon.aws.sensors.sagemaker import SageMakerPipelineSensor + +DESCRIBE_PIPELINE_EXECUTION_RESPONSE = { + "PipelineArn": "string", + "PipelineExecutionArn": "string", + "PipelineExecutionDisplayName": "string", + # Status can be: "Executing" | "Stopping" | "Stopped" | "Failed" | "Succeeded" + "PipelineExecutionStatus": "-- to be set in test --", + "PipelineExecutionDescription": "string", + "PipelineExperimentConfig": {"ExperimentName": "string", "TrialName": "string"}, + "FailureReason": "string", + "CreationTime": datetime(2015, 1, 1), + "LastModifiedTime": datetime(2015, 1, 1), + "CreatedBy": {"UserProfileArn": "string", "UserProfileName": "string", "DomainId": "string"}, + "LastModifiedBy": {"UserProfileArn": "string", "UserProfileName": "string", "DomainId": "string"}, + "ParallelismConfiguration": {"MaxParallelExecutionSteps": 123}, + "ResponseMetadata": { + "HTTPStatusCode": 200, + }, +} + + +class TestSageMakerPipelineSensor: + @staticmethod + def get_response_with_state(state: str): + states = {"Executing", "Stopping", "Stopped", "Failed", "Succeeded"} + assert state in states + res_copy = copy.deepcopy(DESCRIBE_PIPELINE_EXECUTION_RESPONSE) + res_copy["PipelineExecutionStatus"] = state + return res_copy + + @mock.patch.object(SageMakerHook, "get_conn") + @mock.patch.object(SageMakerHook, "describe_pipeline_exec") + def test_sensor_with_failure(self, mock_describe, _): + response_failure = self.get_response_with_state("Failed") + mock_describe.return_value = response_failure + sensor = SageMakerPipelineSensor(pipeline_exec_arn="ARN", task_id="test_task") + + with pytest.raises(AirflowException): + sensor.execute(None) + + mock_describe.assert_called_once_with("ARN", True) + + @mock.patch.object(SageMakerHook, "get_conn") + @mock.patch.object(SageMakerHook, "describe_pipeline_exec") + def test_sensor(self, mock_describe, _): + response_executing = self.get_response_with_state("Executing") + response_stopping = self.get_response_with_state("Stopping") + response_stopped = self.get_response_with_state("Stopped") + mock_describe.side_effect = [ + response_executing, + response_stopping, + response_stopped, + ] + sensor = SageMakerPipelineSensor(pipeline_exec_arn="ARN", task_id="test_task", poke_interval=0) + + sensor.execute(None) + + assert mock_describe.call_count == 3 diff --git a/tests/system/providers/amazon/aws/example_sagemaker.py b/tests/system/providers/amazon/aws/example_sagemaker.py index 994d93fffc266..42a7b9f8a4cbb 100644 --- a/tests/system/providers/amazon/aws/example_sagemaker.py +++ b/tests/system/providers/amazon/aws/example_sagemaker.py @@ -37,11 +37,14 @@ SageMakerDeleteModelOperator, SageMakerModelOperator, SageMakerProcessingOperator, + SageMakerStartPipelineOperator, + SageMakerStopPipelineOperator, SageMakerTrainingOperator, SageMakerTransformOperator, SageMakerTuningOperator, ) from airflow.providers.amazon.aws.sensors.sagemaker import ( + SageMakerPipelineSensor, SageMakerTrainingSensor, SageMakerTransformSensor, SageMakerTuningSensor, @@ -199,6 +202,7 @@ def set_up(env_id, role_arn): training_job_name = f"{env_id}-train" transform_job_name = f"{env_id}-transform" tuning_job_name = f"{env_id}-tune" + pipeline_name = f"{env_id}-pipe" input_data_S3_key = f"{env_id}/processed-input-data" prediction_output_s3_key = f"{env_id}/transform" @@ -218,6 +222,16 @@ def set_up(env_id, role_arn): f"the directions at the top of the system testfile " ) + # Json definition for a dummy pipeline of 30 chained "conditional step" checking that 3 < 6 + # Each step takes roughly 1 second to execute, so the pipeline runtimes is ~30 seconds, which should be + # enough to test stopping and awaiting without race conditions. + # Built using sagemaker sdk, and using json.loads(pipeline.definition()) + pipeline_json_definition = """{"Version": "2020-12-01", "Metadata": {}, "Parameters": [], "PipelineExperimentConfig": {"ExperimentName": {"Get": "Execution.PipelineName"}, "TrialName": {"Get": "Execution.PipelineExecutionId"}}, "Steps": [{"Name": "DummyCond29", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond28", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond27", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond26", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond25", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond24", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond23", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond22", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond21", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond20", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond19", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond18", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond17", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond16", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond15", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond14", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond13", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond12", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond11", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond10", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond9", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond8", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond7", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond6", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond5", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond4", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond3", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond2", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond1", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond0", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [{"Name": "DummyCond", "Type": "Condition", "Arguments": {"Conditions": [{"Type": "LessThanOrEqualTo", "LeftValue": 3.0, "RightValue": 6.0}], "IfSteps": [], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}], "ElseSteps": []}}]}""" # noqa: E501 + sgmk_client = boto3.client("sagemaker") + sgmk_client.create_pipeline( + PipelineName=pipeline_name, PipelineDefinition=pipeline_json_definition, RoleArn=role_arn + ) + resource_config = { "InstanceCount": 1, "InstanceType": "ml.m5.large", @@ -255,7 +269,7 @@ def set_up(env_id, role_arn): "ProcessingResources": { "ClusterConfig": resource_config, }, - "StoppingCondition": {"MaxRuntimeInSeconds": 300}, + "StoppingCondition": {"MaxRuntimeInSeconds": 60}, "AppSpecification": { "ImageUri": ecr_repository_uri, }, @@ -293,7 +307,7 @@ def set_up(env_id, role_arn): "OutputDataConfig": {"S3OutputPath": f"s3://{bucket_name}/{training_output_s3_key}/"}, "ResourceConfig": resource_config, "RoleArn": role_arn, - "StoppingCondition": {"MaxRuntimeInSeconds": 6000}, + "StoppingCondition": {"MaxRuntimeInSeconds": 60}, "TrainingJobName": training_job_name, } model_config = { @@ -354,7 +368,7 @@ def set_up(env_id, role_arn): "OutputDataConfig": {"S3OutputPath": f"s3://{bucket_name}/{training_output_s3_key}"}, "ResourceConfig": resource_config, "RoleArn": role_arn, - "StoppingCondition": {"MaxRuntimeInSeconds": 60000}, + "StoppingCondition": {"MaxRuntimeInSeconds": 60}, }, } transform_config = { @@ -389,6 +403,7 @@ def set_up(env_id, role_arn): ti.xcom_push(key="processing_config", value=processing_config) ti.xcom_push(key="training_config", value=training_config) ti.xcom_push(key="training_job_name", value=training_job_name) + ti.xcom_push(key="pipeline_name", value=pipeline_name) ti.xcom_push(key="model_config", value=model_config) ti.xcom_push(key="model_name", value=model_name) ti.xcom_push(key="tuning_config", value=tuning_config) @@ -421,6 +436,12 @@ def delete_logs(env_id): purge_logs(generated_logs) +@task(trigger_rule=TriggerRule.ALL_DONE) +def delete_pipeline(pipeline_name): + sgmk_client = boto3.client("sagemaker") + sgmk_client.delete_pipeline(PipelineName=pipeline_name) + + with DAG( dag_id=DAG_ID, schedule="@once", @@ -448,6 +469,32 @@ def delete_logs(env_id): replace=True, ) + # [START howto_operator_sagemaker_start_pipeline] + start_pipeline1 = SageMakerStartPipelineOperator( + task_id="start_pipeline1", + pipeline_name=test_setup["pipeline_name"], + ) + # [END howto_operator_sagemaker_start_pipeline] + + # [START howto_operator_sagemaker_stop_pipeline] + stop_pipeline1 = SageMakerStopPipelineOperator( + task_id="stop_pipeline1", + pipeline_exec_arn=start_pipeline1.output, + ) + # [END howto_operator_sagemaker_stop_pipeline] + + start_pipeline2 = SageMakerStartPipelineOperator( + task_id="start_pipeline2", + pipeline_name=test_setup["pipeline_name"], + ) + + # [START howto_sensor_sagemaker_pipeline] + await_pipeline2 = SageMakerPipelineSensor( + task_id="await_pipeline2", + pipeline_exec_arn=start_pipeline2.output, + ) + # [END howto_sensor_sagemaker_pipeline] + # [START howto_operator_sagemaker_processing] preprocess_raw_data = SageMakerProcessingOperator( task_id="preprocess_raw_data", @@ -535,6 +582,10 @@ def delete_logs(env_id): create_bucket, upload_dataset, # TEST BODY + start_pipeline1, + start_pipeline2, + stop_pipeline1, + await_pipeline2, preprocess_raw_data, train_model, await_training, @@ -548,6 +599,7 @@ def delete_logs(env_id): delete_model, delete_bucket, delete_logs(test_context[ENV_ID_KEY]), + delete_pipeline(test_setup["pipeline_name"]), ) from tests.system.utils.watcher import watcher