amoghrajesh commented on code in PR #76036:
URL: https://github.com/apache/airflow/pull/76036#discussion_r4440409754
##########
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)
Review Comment:
Handled in ec733736176
##########
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"]
- 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
- timeout = self.execution_timeout
- if self.max_ingestion_time:
- timeout = datetime.timedelta(seconds=self.max_ingestion_time)
+ def is_job_succeeded(self, status: str) -> bool:
+ return status == "Completed"
- self.defer(
- timeout=timeout,
- trigger=SageMakerTrigger(
- job_name=self.config["TrainingJobName"],
- job_type="Training",
- waiter_delay=self.check_interval,
- waiter_max_attempts=self.max_attempts,
- aws_conn_id=self.aws_conn_id,
- region_name=self.region_name,
- verify=self.verify,
- botocore_config=self.botocore_config,
+ def poll_until_complete(self, external_id: JsonValue, context: Context) ->
None:
Review Comment:
Handled in ec733736176
##########
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:
Handled in ec733736176
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1223,9 +1226,15 @@ 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
Review Comment:
Handled in ec733736176
##########
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:
Handled in ec733736176
##########
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:
Handled in ec733736176
##########
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:
Handled in ec733736176
--
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]
