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 bfe56db36f1 Add on_kill support to SageMakerProcessingOperator (#76413)
bfe56db36f1 is described below
commit bfe56db36f189726c5dcf9468071dbc8ad41f34a
Author: Amogh Desai <[email protected]>
AuthorDate: Mon Nov 9 08:11:13 2026 +0530
Add on_kill support to SageMakerProcessingOperator (#76413)
---
.../providers/amazon/aws/hooks/sagemaker.py | 38 +++++--
.../providers/amazon/aws/operators/sagemaker.py | 36 +++++-
.../providers/amazon/aws/triggers/sagemaker.py | 31 ++++--
.../tests/unit/amazon/aws/hooks/test_sagemaker.py | 36 ++++++
.../aws/operators/test_sagemaker_processing.py | 124 +++++++++++++++++++++
.../aws/operators/test_sagemaker_training.py | 14 +--
.../unit/amazon/aws/triggers/test_sagemaker.py | 19 +++-
7 files changed, 266 insertions(+), 32 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 5f70eaf8f4a..d8f97dda223 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/sagemaker.py
@@ -562,25 +562,47 @@ class SageMakerHook(AwsBaseHook):
)
return response
- def stop_training_job(self, training_job_name: str) -> None:
+ def _stop_job(self, api_name: str, name_kwarg: str, job_kind: str,
job_name: str) -> None:
"""
- Stop a running SageMaker training job.
-
- .. seealso::
- - :external+boto3:py:meth:`SageMaker.Client.stop_training_job`
+ Call a SageMaker ``stop_*`` API, tolerating a job that is no longer
running.
- :param training_job_name: name of the training job to stop.
+ :param api_name: name of the boto3 client method to call.
+ :param name_kwarg: name of the job-name keyword that method takes.
+ :param job_kind: job kind as it should read in the log message.
+ :param job_name: name of the job to stop.
"""
try:
- self.conn.stop_training_job(TrainingJobName=training_job_name)
+ getattr(self.conn, api_name)(**{name_kwarg: 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)
+ self.log.warning("Could not stop %s job %s: %s", job_kind,
job_name, ce)
return
raise
+ 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.
+ """
+ self._stop_job("stop_training_job", "TrainingJobName", "training",
training_job_name)
+
+ def stop_processing_job(self, processing_job_name: str) -> None:
+ """
+ Stop a running SageMaker processing job.
+
+ .. seealso::
+ - :external+boto3:py:meth:`SageMaker.Client.stop_processing_job`
+
+ :param processing_job_name: name of the processing job to stop.
+ """
+ self._stop_job("stop_processing_job", "ProcessingJobName",
"processing", processing_job_name)
+
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 badaeae08e2..4be55c65846 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -272,6 +272,12 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
{"output_name": "evaluation", "file_name": "metrics.json",
"result_name": "EvaluationReport"}
]
# Downstream:
ti.xcom_pull("this_task")["OutputFiles"]["EvaluationReport"]["mse"]
+ :param cancel_on_kill: If True, the processing 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. It does not cover a
deferral that expires on
+ its own ``timeout``: the Triggerer kills the trigger without calling
``on_kill``, and no
+ event reaches the task, so nothing runs to stop the job.
:return: Dict with a ``Processing`` key (the job description) and, when
``output_files_to_xcom``
is configured, an ``OutputFiles`` key containing the parsed JSON
contents of the configured
output files (keyed by their ``result_name``).
@@ -291,6 +297,7 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
action_if_job_exists: str = "timestamp",
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
output_files_to_xcom: list[dict] | None = None,
+ cancel_on_kill: bool = False,
**kwargs,
):
super().__init__(config=config, **kwargs)
@@ -307,7 +314,9 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
self.max_attempts = max_attempts or 60
self.max_ingestion_time = max_ingestion_time
self.deferrable = deferrable
+ self.cancel_on_kill = cancel_on_kill
self.serialized_job: dict
+ self._processing_job_name: str | None = None
def _create_integer_fields(self) -> None:
"""Set fields which should be cast to integers."""
@@ -338,6 +347,9 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
self.hook.describe_processing_job,
)
+ # Recorded before the submit call so a kill in that window still has a
name to stop.
+ self._processing_job_name = self.config["ProcessingJobName"]
+
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."
@@ -383,6 +395,7 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
region_name=self.region_name,
verify=self.verify,
botocore_config=self.botocore_config,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name="execute_complete",
)
@@ -396,6 +409,11 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
validated_event = validate_execute_complete_event(event)
if validated_event["status"] != "success":
+ # A "timeout" event means the waiter gave up while the job was
still running; "failed"
+ # and "stopped" are already terminal.
+ if self.cancel_on_kill and validated_event["status"] == "timeout":
+ self.log.info("Stopping SageMaker processing job %s.",
validated_event["job_name"])
+ self.hook.stop_processing_job(validated_event["job_name"])
raise AirflowException(f"Error while running job:
{validated_event}")
self.log.info("SageMaker job %s completed.",
validated_event["job_name"])
@@ -405,6 +423,12 @@ class SageMakerProcessingOperator(SageMakerBaseOperator):
result.update(self._read_output_files())
return result
+ def on_kill(self) -> None:
+ """Stop the SageMaker processing job if ``cancel_on_kill`` is set."""
+ if self.cancel_on_kill and self._processing_job_name:
+ self.log.info("Stopping SageMaker processing job %s.",
self._processing_job_name)
+ self.hook.stop_processing_job(self._processing_job_name)
+
def _read_output_files(self) -> dict:
"""Read configured output files from S3 after job completion and
return them for XCom."""
if not self.output_files_to_xcom:
@@ -1223,7 +1247,7 @@ 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.
+ :param cancel_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.
@@ -1242,7 +1266,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,
+ cancel_on_kill: bool = False,
**kwargs,
):
super().__init__(config=config, **kwargs)
@@ -1260,7 +1284,7 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
Provided value: '{action_if_job_exists}'."
)
self.deferrable = deferrable
- self.stop_training_job_on_kill = stop_training_job_on_kill
+ self.cancel_on_kill = cancel_on_kill
self.serialized_training_data: dict
self._training_job_name: str | None = None
@@ -1362,7 +1386,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,
+ cancel_on_kill=self.cancel_on_kill,
),
method_name="execute_complete",
)
@@ -1383,8 +1407,8 @@ class SageMakerTrainingOperator(SageMakerBaseOperator):
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:
+ """Stop the SageMaker training job if ``cancel_on_kill`` is set."""
+ if self.cancel_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)
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 e69cb2dfcbb..8dcef92a58a 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
@@ -22,7 +22,7 @@ import warnings
from collections import Counter
from collections.abc import AsyncIterator
from enum import IntEnum
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
from asgiref.sync import sync_to_async
from botocore.exceptions import WaiterError
@@ -54,9 +54,10 @@ 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 cancel_on_kill: If True, stop the job when the user marks the
deferred task failed,
+ clears it, or marks it succeeded. Only the Training and Processing job
types can be stopped;
+ it is a no-op for the others. 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.
"""
@@ -123,17 +124,29 @@ class SageMakerTrigger(AwsBaseWaiterTrigger):
config=self.botocore_config,
)
+ # Spelled via the methods themselves so mypy rejects a name that does not
exist, but stored as
+ # names so the call still resolves on the hook instance.
+ _STOP_METHODS_BY_JOB_TYPE: ClassVar[dict[str, str]] = {
+ "training": SageMakerHook.stop_training_job.__name__,
+ "processing": SageMakerHook.stop_processing_job.__name__,
+ }
+
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":
+ """Stop the job when the user acts on the deferred task."""
+ if not self.cancel_on_kill:
+ return
+ # Absent job types are a no-op, not a call against another type's stop
API.
+ stop_method = self._STOP_METHODS_BY_JOB_TYPE.get(self.job_type.lower())
+ if stop_method is None:
return
- self.log.info("Stopping SageMaker training job %s.", self.job_name)
+ self.log.info("Stopping SageMaker %s job %s.", self.job_type,
self.job_name)
hook: SageMakerHook = self.hook() # type: ignore[assignment]
try:
- await sync_to_async(hook.stop_training_job)(self.job_name)
+ await sync_to_async(getattr(hook, stop_method))(self.job_name)
except Exception:
self.log.exception(
- "Failed to stop SageMaker training job %s. The job may still
be running.",
+ "Failed to stop SageMaker %s job %s. The job may still be
running.",
+ self.job_type,
self.job_name,
)
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 c97f71e13e6..9319a8ff173 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py
@@ -732,6 +732,42 @@ 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_processing_job(self, mock_conn):
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_processing_job(processing_job_name="job_name")
+
+
mock_conn().stop_processing_job.assert_called_once_with(ProcessingJobName="job_name")
+
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_processing_job_swallows_validation_exception(self,
mock_conn):
+ mock_conn().stop_processing_job.side_effect = ClientError(
+ error_response={"Error": {"Code": "ValidationException",
"Message": "already Completed"}},
+ operation_name="StopProcessingJob",
+ )
+
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ hook.stop_processing_job(processing_job_name="job_name")
+
+
mock_conn().stop_processing_job.assert_called_once_with(ProcessingJobName="job_name")
+
+ @patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.conn",
new_callable=mock.PropertyMock)
+ def test_stop_processing_job_raises_other_client_errors(self, mock_conn):
+ error = ClientError(
+ error_response={"Error": {"Code": "AccessDeniedException"}},
+ operation_name="StopProcessingJob",
+ )
+ mock_conn().stop_processing_job.side_effect = error
+
+ hook = SageMakerHook(aws_conn_id="aws_default")
+
+ with pytest.raises(ClientError) as raised_exception:
+ hook.stop_processing_job(processing_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_training_job(self, mock_conn):
hook = SageMakerHook(aws_conn_id="aws_default")
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
index 6ce11d7a669..ab164c3a86f 100644
---
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
+++
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_processing.py
@@ -112,6 +112,130 @@ class TestSageMakerProcessingOperator:
task_id="test_sagemaker_operator", wait_for_completion=True,
check_interval=5, deferrable=True
)
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ def test_on_kill_does_nothing_by_default(self, mock_stop):
+ op = SageMakerProcessingOperator(**self.processing_config_kwargs,
config=CREATE_PROCESSING_PARAMS)
+ op._processing_job_name = "job_name"
+
+ op.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ def test_on_kill_does_nothing_before_the_job_is_submitted(self, mock_stop):
+ op = SageMakerProcessingOperator(
+ **self.processing_config_kwargs,
+ config=CREATE_PROCESSING_PARAMS,
+ cancel_on_kill=True,
+ )
+
+ op.on_kill()
+
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name",
return_value=0)
+ @mock.patch.object(SageMakerHook, "create_processing_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_on_kill_stops_a_job_killed_during_submit(self, _, mock_create,
__, mock_desc, mock_stop):
+ # Only holds while the name is recorded before the submit call.
+ mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
+ op = SageMakerProcessingOperator(
+ task_id="test_sagemaker_operator",
+ config=CREATE_PROCESSING_PARAMS,
+ wait_for_completion=True,
+ check_interval=5,
+ cancel_on_kill=True,
+ )
+ mock_create.side_effect = lambda *args, **kwargs: (
+ op.on_kill(),
+ {"ProcessingJobArn": "test_arn", "ResponseMetadata":
{"HTTPStatusCode": 200}},
+ )[1]
+
+ op.execute(None)
+
+ mock_stop.assert_called_once_with("job_name")
+
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name",
return_value=0)
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={"ProcessingJobArn": "test_arn", "ResponseMetadata":
{"HTTPStatusCode": 200}},
+ )
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_on_kill_stops_the_renamed_job(self, _, mock_create, __,
mock_desc, mock_stop):
+ # The name is rewritten with a timestamp suffix when one already
exists; on_kill has to stop
+ # the name that was actually submitted, not the one in the config.
+ mock_desc.side_effect = [None, ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
+ op = SageMakerProcessingOperator(
+ **self.processing_config_kwargs,
+ config=CREATE_PROCESSING_PARAMS,
+ cancel_on_kill=True,
+ )
+
+ op.execute(None)
+ op.on_kill()
+
+ submitted_name = op.config["ProcessingJobName"]
+ assert submitted_name != "job_name"
+ mock_stop.assert_called_once_with(submitted_name)
+
+ @pytest.mark.parametrize(
+ ("status", "stops"),
+ [("timeout", True), ("failed", False), ("stopped", False)],
+ )
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ def test_execute_complete_stops_a_job_left_running_by_a_timeout(self,
mock_stop, status, stops):
+ op = SageMakerProcessingOperator(
+ **self.processing_config_kwargs,
+ config=CREATE_PROCESSING_PARAMS,
+ cancel_on_kill=True,
+ )
+
+ with pytest.raises(AirflowException):
+ op.execute_complete(None, {"status": status, "job_name":
"job_name"})
+
+ if stops:
+ mock_stop.assert_called_once_with("job_name")
+ else:
+ mock_stop.assert_not_called()
+
+ @mock.patch.object(SageMakerHook, "stop_processing_job")
+ def test_execute_complete_leaves_the_job_alone_when_the_flag_is_off(self,
mock_stop):
+ op = SageMakerProcessingOperator(**self.processing_config_kwargs,
config=CREATE_PROCESSING_PARAMS)
+
+ with pytest.raises(AirflowException):
+ op.execute_complete(None, {"status": "timeout", "job_name":
"job_name"})
+
+ mock_stop.assert_not_called()
+
+ @pytest.mark.parametrize("cancel_on_kill", [True, False])
+ @mock.patch.object(SageMakerHook, "describe_processing_job")
+ @mock.patch.object(SageMakerHook, "count_processing_jobs_by_name",
return_value=0)
+ @mock.patch.object(
+ SageMakerHook,
+ "create_processing_job",
+ return_value={"ProcessingJobArn": "test_arn", "ResponseMetadata":
{"HTTPStatusCode": 200}},
+ )
+ def test_deferred_trigger_receives_stop_on_kill(self, mock_create, __,
mock_desc, cancel_on_kill):
+ mock_desc.side_effect = [
+ ClientError({"Error": {"Code": "ValidationException"}}, "op"),
+ {"ProcessingJobStatus": "InProgress"},
+ ]
+ op = SageMakerProcessingOperator(
+ **self.defer_processing_config_kwargs,
+ config=CREATE_PROCESSING_PARAMS,
+ cancel_on_kill=cancel_on_kill,
+ )
+
+ with pytest.raises(TaskDeferred) as exc:
+ op.execute(None)
+
+ assert exc.value.trigger.cancel_on_kill is cancel_on_kill
+
@mock.patch.object(SageMakerHook, "describe_processing_job")
@mock.patch.object(SageMakerHook, "count_processing_jobs_by_name",
return_value=0)
@mock.patch.object(
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 badd8e798e5..0e309b03c12 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
@@ -126,7 +126,7 @@ class TestSageMakerTrainingOperator:
@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.cancel_on_kill = True
self.sagemaker.on_kill()
@@ -146,7 +146,7 @@ class TestSageMakerTrainingOperator:
check_interval=5,
check_if_job_exists=False,
print_log=False,
- stop_training_job_on_kill=True,
+ cancel_on_kill=True,
)
mock_training.side_effect = lambda *args, **kwargs: (
sagemaker_op.on_kill(),
@@ -169,7 +169,7 @@ class TestSageMakerTrainingOperator:
"TrainingJobArn": "test_arn",
"ResponseMetadata": {"HTTPStatusCode": 200},
}
- self.sagemaker.stop_training_job_on_kill = True
+ self.sagemaker.cancel_on_kill = True
self.sagemaker.execute(None)
self.sagemaker.on_kill()
@@ -308,11 +308,11 @@ 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])
+ @pytest.mark.parametrize("cancel_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
+ self, mock_training, mock_describe_training_job, cancel_on_kill
):
mock_training.return_value = {
"TrainingJobArn": "test_arn",
@@ -325,13 +325,13 @@ class TestSageMakerTrainingOperator:
check_if_job_exists=False,
print_log=False,
deferrable=True,
- stop_training_job_on_kill=stop_training_job_on_kill,
+ cancel_on_kill=cancel_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
+ assert exc.value.trigger.cancel_on_kill is cancel_on_kill
@mock.patch.object(
SageMakerHook,
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 b540741763a..dbd5a8ec4bc 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
@@ -95,8 +95,22 @@ class TestSagemakerTrigger:
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):
+ @pytest.mark.parametrize(
+ ("job_type", "stop_method"),
+ [("Training", "stop_training_job"), ("Processing",
"stop_processing_job")],
+ )
+ async def test_on_kill_stops_the_job_for_its_type(self, job_type,
stop_method):
+ 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()
+
+ getattr(mock_hook, stop_method).assert_called_once_with(JOB_NAME)
+
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize("job_type", ["Transform", "Tuning", "Endpoint"])
+ async def test_on_kill_skips_job_types_without_stop_support(self,
job_type):
trigger = SageMakerTrigger(job_name=JOB_NAME, job_type=job_type,
cancel_on_kill=True)
mock_hook = mock.MagicMock(spec=SageMakerHook)
@@ -104,6 +118,7 @@ class TestSagemakerTrigger:
await trigger.on_kill()
mock_hook.stop_training_job.assert_not_called()
+ mock_hook.stop_processing_job.assert_not_called()
@pytest.mark.asyncio
async def test_on_kill_swallows_stop_training_job_error(self):
