kaxil commented on code in PR #76036:
URL: https://github.com/apache/airflow/pull/76036#discussion_r4429650965
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1274,89 +1283,133 @@ def _create_integer_fields(self) -> None:
def execute(self, context: Context) -> dict:
self.preprocess_config()
+ 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."
+ )
+
+ if self.deferrable and self.wait_for_completion:
+ return self._execute_deferrable(context)
+
+ return self.execute_resumable(context)
+
+ def submit_job(self, context: Context) -> str:
if self.check_if_job_exists:
self.config["TrainingJobName"] = self._get_unique_job_name(
self.config["TrainingJobName"],
self.action_if_job_exists == "fail",
self.hook.describe_training_job,
)
- self.log.info("Creating SageMaker training job %s.",
self.config["TrainingJobName"])
-
- 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.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:
+ return
self.hook.describe_training_job(str(external_id))["TrainingJobStatus"]
Review Comment:
If SageMaker doesn't know the stored name (say `region_name` or
`aws_conn_id` changed and the Dag run was cleared, which keeps task state),
`describe_training_job` raises a `ValidationException` here and every attempt
fails before `submit_job` runs. Glue and Redshift map that to a `NOT_FOUND`
status so the mixin resubmits. Worth doing the same?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1274,89 +1283,133 @@ def _create_integer_fields(self) -> None:
def execute(self, context: Context) -> dict:
self.preprocess_config()
Review Comment:
`preprocess_config()` still runs before the resume decision, and that's
where the `S3Operations` side effects happen. A non-tar `S3Upload` goes through
`load_file` with `replace=False`, so the retry raises `The key ... already
exists` before it ever looks at the stored job, and a tar upload overwrites the
inputs of a job that's still running. Could this move into `submit_job`? The
deferrable path goes through `submit_job` too.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -79,42 +79,127 @@ def setup_method(self):
check_interval=5,
)
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log")
@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")
@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()
+ 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")
+ @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()
+
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log")
+ @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, __):
Review Comment:
These durable tests fail on all the pre-3.3 compat jobs (2.11.1, 3.0.6,
3.1.8, 3.2.2). The compat `ResumableJobMixin` always submits and never touches
`task_state_store`, so `store.set` is never called, and the reconnect test
times out. They need `@pytest.mark.skipif(not AIRFLOW_V_3_3_PLUS, ...)` like
the durable tests in `test_glue.py`.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_training.py:
##########
@@ -79,42 +79,127 @@ def setup_method(self):
check_interval=5,
)
+ @mock.patch.object(SageMakerHook, "check_training_status_with_log")
@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")
@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()
Review Comment:
Could this be `mock.MagicMock(spec_set=["get", "set"])` like
`test_redshift_data.py`? And `autospec=True` on the new
`check_training_status_with_log` patches, so a wrong call signature fails the
test.
--
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]
