pankajkoti commented on code in PR #36685:
URL: https://github.com/apache/airflow/pull/36685#discussion_r1458515289
##########
airflow/providers/amazon/aws/hooks/sagemaker.py:
##########
@@ -1300,3 +1300,13 @@ def create_auto_ml_job(
if "BestCandidate" in res:
return res["BestCandidate"]
return None
+
+ async def describe_training_job_async(self, job_name: str) -> dict[str,
Any]:
Review Comment:
Where does this get used?
##########
airflow/providers/amazon/aws/operators/sagemaker.py:
##########
@@ -1085,8 +1089,50 @@ def execute(self, context: Context) -> dict:
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"]
+
+ 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
+ _, last_description, last_describe_job_call =
self.hook.describe_training_job_with_log(
+ self.config["TrainingJobName"],
+ {},
+ [],
+ instance_count,
+ LogState.TAILING if job_already_completed else
LogState.COMPLETE,
Review Comment:
Can we recheck if the conditional is correct here?
--
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]
