amoghrajesh commented on code in PR #76036:
URL: https://github.com/apache/airflow/pull/76036#discussion_r4520577631
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -79,42 +80,184 @@ def setup_method(self):
check_interval=5,
)
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
@mock.patch.object(SageMakerHook, "describe_training_job")
@mock.patch.object(SageMakerHook, "create_training_job")
@mock.patch.object(sagemaker, "serialize", return_value="")
- def test_integer_fields(self, _, mock_training, mock_desc):
+ def test_integer_fields(self, _, mock_training, mock_desc, __):
mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
mock_training.return_value = {
"TrainingJobArn": "test_arn",
"ResponseMetadata": {"HTTPStatusCode": 200},
}
- self.sagemaker.execute(None)
+ self.sagemaker.execute({})
assert self.sagemaker.integer_fields == EXPECTED_INTEGER_FIELDS
for key1, key2 in EXPECTED_INTEGER_FIELDS:
assert self.sagemaker.config[key1][key2] ==
int(self.sagemaker.config[key1][key2])
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
@mock.patch.object(SageMakerHook, "describe_training_job")
@mock.patch.object(SageMakerHook, "create_training_job")
@mock.patch.object(sagemaker, "serialize", return_value="")
- def test_execute_without_check_if_job_exists(self, _, mock_training,
mock_desc):
+ def test_execute_without_check_if_job_exists(self, _, mock_training,
mock_desc, __):
mock_training.return_value = {
"TrainingJobArn": "test_arn",
"ResponseMetadata": {"HTTPStatusCode": 200},
}
self.sagemaker.check_if_job_exists = False
self.sagemaker._check_if_job_exists = mock.MagicMock()
- self.sagemaker.execute(None)
+ self.sagemaker.execute({})
self.sagemaker._check_if_job_exists.assert_not_called()
mock_training.assert_called_once_with(
CREATE_TRAINING_PARAMS,
wait_for_completion=False,
- print_log=True,
+ print_log=False,
check_interval=5,
max_ingestion_time=None,
)
- a = []
- a.sort()
+ @staticmethod
+ def _task_state_store(stored_job_name=None):
+ store = mock.MagicMock(spec_set=["get", "set"])
+ store.get.return_value = stored_job_name
+ return store
+
+ @pytest.mark.parametrize(
+ ("status", "active", "succeeded"),
+ [
+ ("InProgress", True, False),
+ ("Stopping", True, False),
+ ("Completed", False, True),
+ ("Failed", False, False),
+ ("Stopped", False, False),
+ ],
+ )
+ def test_job_status_mapping(self, status, active, succeeded):
+ assert self.sagemaker.is_job_active(status) is active
+ assert self.sagemaker.is_job_succeeded(status) is succeeded
+
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_durable_disabled_ignores_the_stored_job(self, _, mock_training,
mock_desc, __):
+ mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ self.sagemaker.durable = False
+ store = self._task_state_store("previously_running_job")
+
+ self.sagemaker.execute({"task_state_store": store})
+
+ mock_training.assert_called_once()
+ store.get.assert_not_called()
+ store.set.assert_not_called()
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS,
+ reason="ResumableJobMixin reconnect requires task_state_store,
available in Airflow 3.3+",
+ )
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_durable_persists_the_submitted_job_name(self, _, mock_training,
mock_desc, __):
+ mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), None]
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ store = self._task_state_store()
+
+ self.sagemaker.execute({"task_state_store": store})
+
+ store.set.assert_called_once_with("sagemaker_training_job_name",
"job_name")
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS,
+ reason="ResumableJobMixin reconnect requires task_state_store,
available in Airflow 3.3+",
+ )
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ @mock.patch.object(SageMakerHook, "configure_s3_resources")
+ def test_durable_reconnects_to_a_running_job(
+ self, mock_configure_s3, _, mock_training, mock_desc, mock_poll
+ ):
+ mock_desc.return_value = {"TrainingJobStatus": "InProgress"}
+ store = self._task_state_store("previously_running_job")
+
+ result = self.sagemaker.execute({"task_state_store": store})
+
+ mock_training.assert_not_called()
+ mock_configure_s3.assert_not_called()
+ # autospec leaves the hook instance as args[0], so the job name is
args[1].
+ assert mock_poll.call_args.args[1] == "previously_running_job"
+ assert result == {"Training": ""}
+ assert self.sagemaker.serialized_training_data == ""
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS,
+ reason="ResumableJobMixin reconnect requires task_state_store,
available in Airflow 3.3+",
+ )
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_durable_resubmits_when_the_stored_job_is_unknown(self, _,
mock_training, mock_desc, __):
+ # The stored name is unresolvable after a region or connection change.
+ # Only the status lookup for the stored name fails; the describe
behind serialize_result works.
+ mock_desc.side_effect = [ClientError({"Error": {"Code":
"ValidationException"}}, "op"), {}]
+ mock_training.return_value = {
+ "TrainingJobArn": "test_arn",
+ "ResponseMetadata": {"HTTPStatusCode": 200},
+ }
+ self.sagemaker.check_if_job_exists = False
+ store = self._task_state_store("job_from_another_account")
+
+ self.sagemaker.execute({"task_state_store": store})
+
+ mock_training.assert_called_once()
+ store.set.assert_called_once_with("sagemaker_training_job_name",
"job_name")
+
+ def test_unknown_job_status_is_neither_active_nor_succeeded(self):
+ assert self.sagemaker.is_job_active("NOT_FOUND") is False
+ assert self.sagemaker.is_job_succeeded("NOT_FOUND") is False
+
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ def test_get_job_status_reraises_other_client_errors(self, mock_desc):
+ error = ClientError({"Error": {"Code": "AccessDeniedException"}}, "op")
+ mock_desc.side_effect = error
+
+ with pytest.raises(ClientError) as raised:
+ self.sagemaker.get_job_status("job_name", {})
+
+ assert raised.value is error
+
+ @pytest.mark.skipif(
+ not AIRFLOW_V_3_3_PLUS,
+ reason="ResumableJobMixin reconnect requires task_state_store,
available in Airflow 3.3+",
+ )
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log",
autospec=True)
+ @mock.patch.object(SageMakerHook, "describe_training_job")
+ @mock.patch.object(SageMakerHook, "create_training_job")
+ @mock.patch.object(sagemaker, "serialize", return_value="")
+ def test_durable_resubmits_after_a_terminal_job(self, _, mock_training,
mock_desc, __):
Review Comment:
Added. Confirmed it is the only test that fails if the `_result` assignment
goes.
--
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]
