kaxil commented on code in PR #76033:
URL: https://github.com/apache/airflow/pull/76033#discussion_r4425091460
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -116,6 +116,52 @@ def test_execute_without_check_if_job_exists(self, _,
mock_training, mock_desc):
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")
+ def test_on_kill_stops_the_job_when_enabled(self, mock_stop):
+ self.sagemaker.stop_training_job_on_kill = True
+ self.sagemaker._training_job_name = "job_name"
+
+ self.sagemaker.on_kill()
+
+ 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):
Review Comment:
None of these tests would fail if the `_training_job_name` assignment moved
below `create_training_job`: they either set it by hand or call `on_kill` after
`execute()` returns, and the fixture uses `wait_for_completion=False`. A test
with `wait_for_completion=True` and a `mock_training.side_effect` that calls
`self.sagemaker.on_kill()` would pin the ordering the PR relies on. Passing
`stop_training_job_on_kill=True` through the constructor instead of setting the
attribute would also cover the `__init__` wiring, and that test could replace
`test_on_kill_stops_the_job_when_enabled`, which hits the same branch as this
one.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1371,6 +1379,12 @@ def serialize_result(self, job_name: str) -> dict[str,
dict]:
self.serialized_training_data =
serialize(self.hook.describe_training_job(job_name))
return {"Training": self.serialized_training_data}
+ def on_kill(self) -> None:
Review Comment:
With `deferrable=True` this doesn't run for most of the job's life. Once
`execute` defers to `SageMakerTrigger`, mark-failed or clear goes to the
trigger's `on_kill`, and `SageMakerTrigger` doesn't override it, so a deferred
training job keeps running even with the flag set. Long training jobs are
exactly the ones people run deferred, and `default_deferrable=True` turns that
on deployment-wide.
Could we pass the flag into `SageMakerTrigger` and add an `async on_kill`
there, the way `EmrContainerTrigger` does with `cancel_on_kill`? The trigger is
shared with the other SageMaker operators, so defaulting it to `False` keeps
them unchanged. Trigger-side `on_kill` only fires on Airflow 3.3+, which the
docstring should mention. If that's out of scope here, the docstring at least
needs to say the flag only applies when the task isn't deferred.
##########
providers/amazon/tests/unit/amazon/aws/hooks/test_sagemaker.py:
##########
@@ -729,6 +729,40 @@ def test_start_pipeline_returns_arn(self, mock_conn):
# 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")
Review Comment:
This only checks that nothing raised, so it would still pass if the method
returned before calling the API. Could it also assert
`mock_conn().stop_training_job.assert_called_once_with(TrainingJobName="job_name")`
and the warning (via `caplog`)?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
