ramitkataria commented on code in PR #71653:
URL: https://github.com/apache/airflow/pull/71653#discussion_r4023275737
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py:
##########
@@ -108,6 +109,41 @@ def hook(self) -> AwsGenericHook:
config=self.botocore_config,
)
+ def _event_from_exception(self, error: AirflowException) -> TriggerEvent:
+
+ if isinstance(error, WaiterMaxAttemptsError):
Review Comment:
Fair, let's keep that out of scope for this PR. Might be worth a follow-up
in `async_wait` to raise a distinct error when every attempt failed on
credentials, so triggers can report it as `error` instead of `timeout`.
##########
providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py:
##########
@@ -121,6 +123,91 @@ async def test_sagemaker_trigger_run_all_job_types(self,
mock_async_conn, mock_g
assert response == TriggerEvent({"status": "success", "job_name":
JOB_NAME})
+ @pytest.mark.parametrize(
+ ("aws_status", "expected_status"),
+ [
+ pytest.param("Failed", "failed", id="failed"),
+ pytest.param("Stopped", "stopped", id="stopped"),
+ ],
+ )
+ def test_event_from_exception_terminal_state(self, aws_status,
expected_status):
+ trigger = SageMakerTrigger(
+ job_name=JOB_NAME,
+ job_type=JOB_TYPE,
+ waiter_delay=WAITER_DELAY,
+ waiter_max_attempts=WAITER_MAX_ATTEMPTS,
+ aws_conn_id=AWS_CONN_ID,
+ )
+ error = WaiterTerminalFailure(
+ "SageMaker job failed",
+ last_response={"TrainingJobStatus": aws_status},
Review Comment:
Should we add a parametrize case for the `FailureReason` path?
--
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]
