kaxil commented on code in PR #76036:
URL: https://github.com/apache/airflow/pull/76036#discussion_r4515728412
##########
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:
Nothing covers a stored job that's already `Completed`, which is the only
path where `get_job_result` fills `_result` itself. If that assignment goes,
`execute()` returns `None` and every test still passes. Could you add one next
to this that asserts no submit, no `store.set`, and the returned dict?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1223,13 +1226,19 @@ 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 durable: If True (the default), the training job name is persisted
to the task state store
+ before polling, and a retry reconnects to that job instead of
submitting a new one. Only applies
+ to the synchronous path; a deferrable run is unaffected. Needs Airflow
3.3 or later, below it the
+ operator behaves as with ``durable=False``.
: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.
"""
+ external_id_key = "sagemaker_training_job_name"
+
def __init__(
Review Comment:
Glue and Redshift Data declare `durable: bool | None = None` and forward it
only when it's set. Here it stays in `**kwargs`, and below 3.3 the compat
stand-in's `__init__` isn't wrapped by `BaseOperatorMeta`, so
`SageMakerTrainingOperator.partial(durable=False)` fails parsing with
`unexpected keyword argument 'durable'`. Could this follow the same pattern?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1289,85 +1315,109 @@ def execute(self, context: Context) -> dict:
)
# 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(
- "Setting deferrable to True does not have effect when
wait_for_completion is set to False."
- )
-
- wait_for_completion = self.wait_for_completion
- if self.deferrable and self.wait_for_completion:
- # Set wait_for_completion to False so that it waits for the status
in the deferred task.
- wait_for_completion = False
+ job_name = self._training_job_name = self.config["TrainingJobName"]
+ self.log.info("Creating SageMaker training job %s.", job_name)
response = self.hook.create_training_job(
self.config,
- wait_for_completion=wait_for_completion,
- print_log=self.print_log,
+ wait_for_completion=False,
+ print_log=False,
check_interval=self.check_interval,
max_ingestion_time=self.max_ingestion_time,
)
if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
raise AirflowException(f"Sagemaker Training Job creation failed:
{response}")
- if self.deferrable and self.wait_for_completion:
- description =
self.hook.describe_training_job(self.config["TrainingJobName"])
- status = description["TrainingJobStatus"]
+ return job_name
- if self.print_log:
- instance_count = description["ResourceConfig"]["InstanceCount"]
- last_describe_job_call = time.monotonic()
- job_already_completed = status not in
self.hook.non_terminal_states
- _, description, last_describe_job_call =
self.hook.describe_training_job_with_log(
- self.config["TrainingJobName"],
- {},
- [],
- instance_count,
- LogState.COMPLETE if job_already_completed else
LogState.TAILING,
- description,
- last_describe_job_call,
- )
- self.log.info(secondary_training_status_message(description,
None))
+ def get_job_status(self, external_id: JsonValue, context: Context) -> str:
+ """Query the raw job status; a job name SageMaker no longer knows
degrades to NOT_FOUND."""
+ try:
+ return
self.hook.describe_training_job(str(external_id))["TrainingJobStatus"]
+ except ClientError as e:
+ if e.response["Error"]["Code"] in ("ValidationException",
"ResourceNotFound"):
+ return "NOT_FOUND"
+ raise
- if status in self.hook.failed_states:
- reason = description.get("FailureReason", "(No reason
provided)")
- raise AirflowException(f"SageMaker job failed because
{reason}")
- if status == "Completed":
- log_message = f"{self.task_id} completed successfully."
- if self.print_log:
- billable_seconds = SageMakerHook.count_billable_seconds(
- training_start_time=description["TrainingStartTime"],
- training_end_time=description["TrainingEndTime"],
- instance_count=instance_count,
- )
- log_message = f"Billable seconds:
{billable_seconds}\n{log_message}"
- self.log.info(log_message)
- return {"Training": serialize(description)}
+ def is_job_active(self, status: str) -> bool:
+ return status in self.hook.non_terminal_states
Review Comment:
If the stored job is `Stopping` (say `on_kill` with
`stop_training_job_on_kill=True` stopped it when the last try timed out), this
reconnects to it. With `print_log=False`, `check_status` treats `Stopped` as
done, so the task succeeds on a job that never finished. Could this be `status
== "InProgress"` so a stopping job gets a fresh submit instead?
--
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]
