This is an automated email from the ASF dual-hosted git repository.
amoghrajesh pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fdeff70fc0b Add on_kill support to SageMakerTrainingOperator (#76033)
fdeff70fc0b is described below
commit fdeff70fc0b879cd1a45c2f0bf65cb55a7856750
Author: Amogh Desai <[email protected]>
AuthorDate: Tue Nov 3 18:43:37 2026 +0530
Add on_kill support to SageMakerTrainingOperator (#76033)
---
.../providers/amazon/aws/hooks/sagemaker.py | 19 +++++
.../providers/amazon/aws/operators/sagemaker.py | 19 ++++-
.../providers/amazon/aws/triggers/sagemaker.py | 26 ++++++-
.../tests/unit/amazon/aws/hooks/test_sagemaker.py | 59 +++++++++++++++
.../aws/operators/test_sagemaker_training.py | 87 ++++++++++++++++++++++
.../unit/amazon/aws/triggers/test_sagemaker.py | 55 ++++++++++++++
6 files changed, 263 insertions(+), 2 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
index f071ff844c9..5f70eaf8f4a 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
@@ -562,6 +562,25 @@ class SageMakerHook(AwsBaseHook):
)
return response
+ def stop_training_job(self, training_job_name: str) -> None:
+ """
+ Stop a running SageMaker training job.
+
+ .. seealso::
+ - :external+boto3:py:meth:`SageMaker.Client.stop_training_job`
+
+ :param training_job_name: name of the training job to stop.
+ """
+ try:
+ self.conn.stop_training_job(TrainingJobName=training_job_name)
+ except ClientError as ce:
+ # SageMaker only accepts a stop for a job in InProgress; a
ValidationException here means
+ # the job reached a terminal state (or never existed) before the
stop landed.
+ if ce.response["Error"]["Code"] == "ValidationException":
+ self.log.warning("Could not stop training job %s: %s",
training_job_name, ce)
+ return
+ raise
+
def describe_training_job(self, name: str):
"""
Get the training job info associated with the name.
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
index 661c58fed14..badaeae08e2 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -1223,6 +1223,10 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
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.
+ :param stop_training_job_on_kill: If True, the training job is stopped
when the task is killed.
+ Without it the job keeps running, and billing, after the task is gone.
A deferred task is
+ stopped from the Triggerer, which needs Airflow 3.3 or later; below
that the flag only
+ applies while the task holds a worker slot.
:return Dict: Returns The ARN of the training job created in Amazon
SageMaker.
"""
@@ -1238,6 +1242,7 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
check_if_job_exists: bool = True,
action_if_job_exists: str = "timestamp",
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
+ stop_training_job_on_kill: bool = False,
**kwargs,
):
super().__init__(config=config, **kwargs)
@@ -1255,7 +1260,9 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
+ self.stop_training_job_on_kill = stop_training_job_on_kill
self.serialized_training_data: dict
+ self._training_job_name: str | None = None
def expand_role(self) -> None:
"""Expand an IAM role name into an ARN."""
@@ -1281,7 +1288,10 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
self.hook.describe_training_job,
)
- self.log.info("Creating SageMaker training job %s.",
self.config["TrainingJobName"])
+ # Recorded before the submit call so a kill in that window still has a
name to stop.
+ self._training_job_name = self.config["TrainingJobName"]
+
+ self.log.info("Creating SageMaker training job %s.",
self._training_job_name)
if self.deferrable and not self.wait_for_completion:
self.log.warning(
@@ -1352,6 +1362,7 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
region_name=self.region_name,
verify=self.verify,
botocore_config=self.botocore_config,
+ cancel_on_kill=self.stop_training_job_on_kill,
),
method_name="execute_complete",
)
@@ -1371,6 +1382,12 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
self.serialized_training_data =
serialize(self.hook.describe_training_job(job_name))
return {"Training": self.serialized_training_data}
+ def on_kill(self) -> None:
+ """Stop the SageMaker training job if ``stop_training_job_on_kill`` is
set."""
+ if self.stop_training_job_on_kill and self._training_job_name:
+ self.log.info("Stopping SageMaker training job %s.",
self._training_job_name)
+ self.hook.stop_training_job(self._training_job_name)
+
def get_openlineage_facets_on_complete(self, task_instance) ->
OperatorLineage:
"""Return OpenLineage data gathered from SageMaker's API response
saved by training job."""
from airflow.providers.openlineage.extractors import OperatorLineage
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
index cc4725e8257..e69cb2dfcbb 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
@@ -24,6 +24,7 @@ from collections.abc import AsyncIterator
from enum import IntEnum
from typing import TYPE_CHECKING
+from asgiref.sync import sync_to_async
from botocore.exceptions import WaiterError
from airflow.exceptions import AirflowProviderDeprecationWarning
@@ -53,6 +54,9 @@ class SageMakerTrigger(AwsBaseWaiterTrigger):
:param region_name: The AWS region where the job is running. Used to build
the hook.
:param verify: Whether or not to verify SSL certificates. Used to build
the hook.
:param botocore_config: Configuration dictionary for the botocore client.
Used to build the hook.
+ :param cancel_on_kill: If True, stop the training job when the user marks
the deferred task failed,
+ clears it, or marks it succeeded. Only applies when ``job_type`` is
``Training``. Requires
+ ``apache-airflow`` with ``BaseTrigger.on_kill()`` support (3.3 and
later); below that it is inert.
:param poke_interval: (deprecated) use ``waiter_delay`` instead.
:param max_attempts: (deprecated) use ``waiter_max_attempts`` instead.
"""
@@ -67,6 +71,7 @@ class SageMakerTrigger(AwsBaseWaiterTrigger):
region_name: str | None = None,
verify: bool | str | None = None,
botocore_config: dict | None = None,
+ cancel_on_kill: bool = False,
poke_interval: int | None = None,
max_attempts: int | None = None,
):
@@ -88,8 +93,13 @@ class SageMakerTrigger(AwsBaseWaiterTrigger):
waiter_max_attempts = max_attempts
self.job_name = job_name
self.job_type = job_type
+ self.cancel_on_kill = cancel_on_kill
super().__init__(
- serialized_fields={"job_name": job_name, "job_type": job_type},
+ serialized_fields={
+ "job_name": job_name,
+ "job_type": job_type,
+ "cancel_on_kill": cancel_on_kill,
+ },
waiter_name=self._get_job_type_waiter(job_type),
waiter_args={self._get_waiter_arg_name(job_type): job_name},
failure_message=f"Error while waiting for {job_type} job",
@@ -113,6 +123,20 @@ class SageMakerTrigger(AwsBaseWaiterTrigger):
config=self.botocore_config,
)
+ async def on_kill(self) -> None:
+ """Stop the training job when the user acts on the deferred task."""
+ if not self.cancel_on_kill or self.job_type.lower() != "training":
+ return
+ self.log.info("Stopping SageMaker training job %s.", self.job_name)
+ hook: SageMakerHook = self.hook() # type: ignore[assignment]
+ try:
+ await sync_to_async(hook.stop_training_job)(self.job_name)
+ except Exception:
+ self.log.exception(
+ "Failed to stop SageMaker training job %s. The job may still
be running.",
+ self.job_name,
+ )
+
def _event_from_exception(self, error: AirflowException) -> TriggerEvent:
if isinstance(error, WaiterNoCredentialsError):
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py
index 4b10fcff20d..c97f71e13e6 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py
@@ -17,6 +17,7 @@
# under the License.
from __future__ import annotations
+import re
import time
from datetime import UTC, datetime
from unittest import mock
@@ -37,6 +38,8 @@ from airflow.providers.amazon.aws.hooks.sagemaker import (
)
from airflow.providers.common.compat.sdk import AirflowException
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_1_PLUS
+
role = "arn:aws:iam:role/test-role"
path = "local/data"
@@ -729,6 +732,62 @@ class TestSageMakerHook:
# 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_stop_training_job(self, mock_conn):
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_training_job(training_job_name="job_name")
+
+
mock_conn().stop_training_job.assert_called_once_with(TrainingJobName="job_name")
+
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_training_job_swallows_validation_exception(self, mock_conn):
+ mock_conn().stop_training_job.side_effect = ClientError(
+ error_response={"Error": {"Code": "ValidationException",
"Message": "already Completed"}},
+ operation_name="StopTrainingJob",
+ )
+
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_training_job(training_job_name="job_name")
+
+
mock_conn().stop_training_job.assert_called_once_with(TrainingJobName="job_name")
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_1_PLUS,
+ reason="caplog only supports structured assertions under structlog,
from Airflow 3.1+",
+ )
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_training_job_warns_when_the_job_is_not_running(self,
mock_conn, caplog):
+ mock_conn().stop_training_job.side_effect = ClientError(
+ error_response={"Error": {"Code": "ValidationException",
"Message": "already Completed"}},
+ operation_name="StopTrainingJob",
+ )
+
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_training_job(training_job_name="job_name")
+
+ assert {
+ "event": re.compile(r"Could not stop training job job_name"),
+ "log_level": "warning",
+ } in caplog
+
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_training_job_raises_other_client_errors(self, mock_conn):
+ error = ClientError(
+ error_response={"Error": {"Code": "AccessDeniedException"}},
+ operation_name="StopTrainingJob",
+ )
+ mock_conn().stop_training_job.side_effect = error
+
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ with pytest.raises(ClientError) as raised_exception:
+ hook.stop_training_job(training_job_name="job_name")
+
+ assert raised_exception.value == error
+
@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"}
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py
index abd1eccb848..badd8e798e5 100644
---
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py
+++
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py
@@ -116,6 +116,68 @@ class TestSageMakerTrainingOperator:
a = []
a.sort()
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ def test_on_kill_does_nothing_by_default(self, mock_stop):
+ self.sagemaker._training_job_name = "job_name"
+
+ self.sagemaker.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ def test_on_kill_does_nothing_before_the_job_is_submitted(self, mock_stop):
+ self.sagemaker.stop_training_job_on_kill = True
+
+ self.sagemaker.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_on_kill_stops_a_job_killed_during_submit(self, _, mock_training,
mock_desc, mock_stop):
+ # A kill landing while create_training_job is in flight must still
find a name to stop, which
+ # only holds while the name is recorded before the submit call.
+ sagemaker_op = SageMakerTrainingOperator(
+ task_id="test_sagemaker_operator",
+ config=CREATE_TRAINING_PARAMS.copy(),
+ wait_for_completion=True,
+ check_interval=5,
+ check_if_job_exists=False,
+ print_log=False,
+ stop_training_job_on_kill=True,
+ )
+ mock_training.side_effect = lambda *args, **kwargs: (
+ sagemaker_op.on_kill(),
+ {"TrainingJobArn": "test_arn", "ResponseMetadata":
{"HTTPStatusCode": 200}},
+ )[1]
+
+ sagemaker_op.execute(None)
+
+ mock_stop.assert_called_once_with("job_name")
+
+ @mock.patch.object(SageMakerHook, "stop_training_job")
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_on_kill_stops_the_renamed_job(self, _, mock_training, mock_desc,
mock_stop):
+ # The job name is rewritten with a timestamp suffix when one already
exists; on_kill has to
+ # stop the name that was actually submitted.
+ mock_desc.side_effect = [None, ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ self.sagemaker.stop_training_job_on_kill = True
+
+ self.sagemaker.execute(None)
+ self.sagemaker.on_kill()
+
+ submitted_name = self.sagemaker.config["TrainingJobName"]
+ assert submitted_name != "job_name"
+ mock_stop.assert_called_once_with(submitted_name)
+
@mock.patch.object(SageMakerHook, "describe_training_job")
@mock.patch.object(SageMakerHook, "create_training_job")
def test_execute_with_failure(self, mock_training, mock_desc):
@@ -246,6 +308,31 @@ class TestSageMakerTrainingOperator:
assert exc.value.trigger.verify == VERIFY
assert exc.value.trigger.botocore_config == BOTOCORE_CONFIG
+ @pytest.mark.parametrize("stop_training_job_on_kill", [True, False])
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ def test_deferred_trigger_receives_stop_on_kill(
+ self, mock_training, mock_describe_training_job,
stop_training_job_on_kill
+ ):
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ mock_describe_training_job.return_value = {"TrainingJobStatus":
"InProgress"}
+ sagemaker = SageMakerTrainingOperator(
+ task_id="test_sagemaker_operator",
+ config=CREATE_TRAINING_PARAMS.copy(),
+ check_if_job_exists=False,
+ print_log=False,
+ deferrable=True,
+ stop_training_job_on_kill=stop_training_job_on_kill,
+ )
+
+ with pytest.raises(TaskDeferred) as exc:
+ sagemaker.execute(context=None)
+
+ assert exc.value.trigger.cancel_on_kill is stop_training_job_on_kill
+
@mock.patch.object(
SageMakerHook,
"describe_training_job",
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
index 2dde688b16b..b540741763a 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
@@ -60,6 +60,61 @@ class TestSagemakerTrigger:
assert args["waiter_max_attempts"] == WAITER_MAX_ATTEMPTS
assert args["aws_conn_id"] == AWS_CONN_ID
assert args["region_name"] == REGION_NAME
+ assert args["cancel_on_kill"] is False
+
+ @pytest.mark.parametrize("cancel_on_kill", [True, False])
+ def test_sagemaker_trigger_serializes_cancel_on_kill(self, cancel_on_kill):
+ trigger = SageMakerTrigger(
+ job_name=JOB_NAME,
+ job_type=JOB_TYPE,
+ cancel_on_kill=cancel_on_kill,
+ )
+
+ _, args = trigger.serialize()
+
+ assert args["cancel_on_kill"] is cancel_on_kill
+
+ @pytest.mark.asyncio
+ async def test_on_kill_stops_the_training_job(self):
+ trigger = SageMakerTrigger(job_name=JOB_NAME, job_type=JOB_TYPE,
cancel_on_kill=True)
+ mock_hook = mock.MagicMock(spec=SageMakerHook)
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+ mock_hook.stop_training_job.assert_called_once_with(JOB_NAME)
+
+ @pytest.mark.asyncio
+ async def test_on_kill_does_nothing_when_not_enabled(self):
+ trigger = SageMakerTrigger(job_name=JOB_NAME, job_type=JOB_TYPE,
cancel_on_kill=False)
+ mock_hook = mock.MagicMock(spec=SageMakerHook)
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+ mock_hook.stop_training_job.assert_not_called()
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("job_type", ["Transform", "Processing", "Tuning",
"Endpoint"])
+ async def test_on_kill_only_stops_training_jobs(self, job_type):
+ trigger = SageMakerTrigger(job_name=JOB_NAME, job_type=job_type,
cancel_on_kill=True)
+ mock_hook = mock.MagicMock(spec=SageMakerHook)
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+ mock_hook.stop_training_job.assert_not_called()
+
+ @pytest.mark.asyncio
+ async def test_on_kill_swallows_stop_training_job_error(self):
+ trigger = SageMakerTrigger(job_name=JOB_NAME, job_type=JOB_TYPE,
cancel_on_kill=True)
+ mock_hook = mock.MagicMock(spec=SageMakerHook)
+ mock_hook.stop_training_job.side_effect = RuntimeError("boom")
+
+ with mock.patch.object(trigger, "hook", return_value=mock_hook):
+ await trigger.on_kill()
+
+ mock_hook.stop_training_job.assert_called_once_with(JOB_NAME)
@pytest.mark.parametrize(
("deprecated_kwarg", "canonical_attr", "value"),
