amoghrajesh commented on code in PR #76036:
URL: https://github.com/apache/airflow/pull/76036#discussion_r4520576007
##########
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:
Narrowed to `status == "InProgress"`, so a `Stopping` job resubmits. Tests
updated.
##########
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:
Done, same shape as Glue and Redshift Data.
Note: `default_args` / `partial()` only break below 3.3, so the test also
asserts `durable` is in the signature.
--
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]
