This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 763cb20ddc1 Fix SageMaker credential failures being reported as
timeouts (#74867)
763cb20ddc1 is described below
commit 763cb20ddc1ce25475bec64b479675ea0fc03bb2
Author: Mohamed Laghdaf Dadda <[email protected]>
AuthorDate: Fri Oct 23 13:44:27 2026 +0100
Fix SageMaker credential failures being reported as timeouts (#74867)
---
.../providers/amazon/aws/triggers/sagemaker.py | 9 ++-
.../unit/amazon/aws/triggers/test_sagemaker.py | 65 +++++++++++++++++++++-
2 files changed, 71 insertions(+), 3 deletions(-)
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
index cbaeb7a1c04..cc4725e8257 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py
@@ -27,7 +27,11 @@ from typing import TYPE_CHECKING
from botocore.exceptions import WaiterError
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.amazon.aws.exceptions import WaiterMaxAttemptsError,
WaiterTerminalFailure
+from airflow.providers.amazon.aws.exceptions import (
+ WaiterMaxAttemptsError,
+ WaiterNoCredentialsError,
+ WaiterTerminalFailure,
+)
from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
from airflow.providers.common.compat.sdk import AirflowException
@@ -111,6 +115,9 @@ class SageMakerTrigger(AwsBaseWaiterTrigger):
def _event_from_exception(self, error: AirflowException) -> TriggerEvent:
+ if isinstance(error, WaiterNoCredentialsError):
+ return super()._event_from_exception(error)
+
if isinstance(error, WaiterMaxAttemptsError):
return TriggerEvent(
{
diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
b/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
index b736505de3c..2dde688b16b 100644
--- a/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
+++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_sagemaker.py
@@ -20,10 +20,14 @@ from unittest import mock
from unittest.mock import AsyncMock
import pytest
-from botocore.exceptions import WaiterError
+from botocore.exceptions import NoCredentialsError, WaiterError
from airflow.exceptions import AirflowProviderDeprecationWarning
-from airflow.providers.amazon.aws.exceptions import WaiterMaxAttemptsError,
WaiterTerminalFailure
+from airflow.providers.amazon.aws.exceptions import (
+ WaiterMaxAttemptsError,
+ WaiterNoCredentialsError,
+ WaiterTerminalFailure,
+)
from airflow.providers.amazon.aws.hooks.sagemaker import SageMakerHook
from airflow.providers.amazon.aws.triggers.sagemaker import
SageMakerPipelineTrigger, SageMakerTrigger
from airflow.providers.common.compat.sdk import AirflowException
@@ -212,6 +216,10 @@ class TestSagemakerTrigger:
),
id="unknown-status",
),
+ pytest.param(
+ WaiterNoCredentialsError("SageMaker job failed"),
+ id="no-credentials",
+ ),
pytest.param(
AirflowException("SageMaker job failed"),
id="generic-error",
@@ -235,6 +243,59 @@ class TestSagemakerTrigger:
}
)
+ @pytest.mark.asyncio
+ @pytest.mark.parametrize(
+ ("errors", "expected_status", "expected_message"),
+ [
+ pytest.param(
+ [NoCredentialsError(), NoCredentialsError()],
+ "error",
+ "Waiter error: max attempts reached due to missing
credentials: Unable to locate credentials",
+ id="no-credentials",
+ ),
+ pytest.param(
+ [
+ WaiterError("test_waiter", "Max attempts exceeded", {}),
+ WaiterError("test_waiter", "Max attempts exceeded", {}),
+ ],
+ "timeout",
+ "Waiter error: max attempts reached",
+ id="ordinary-exhaustion",
+ ),
+ pytest.param(
+ [WaiterError("test_waiter", "terminal failure",
{"TrainingJobStatus": "Failed"})],
+ "failed",
+ "Error while waiting for training job: Failed\nWaiter
test_waiter failed: terminal failure",
+ id="terminal-failure",
+ ),
+ pytest.param(
+ [AirflowException("Unexpected failure")],
+ "error",
+ "Unexpected failure",
+ id="generic-error",
+ ),
+ ],
+ )
+
@mock.patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.get_waiter")
+
@mock.patch("airflow.providers.amazon.aws.hooks.sagemaker.SageMakerHook.get_async_conn")
+ async def test_run_waiter_errors(
+ self, mock_async_conn, mock_get_waiter, errors, expected_status,
expected_message
+ ):
+ mock_get_waiter.return_value.wait = AsyncMock(side_effect=errors)
+ trigger = SageMakerTrigger(
+ job_name=JOB_NAME,
+ job_type=JOB_TYPE,
+ waiter_delay=0,
+ waiter_max_attempts=2,
+ aws_conn_id=AWS_CONN_ID,
+ )
+
+ events = [event async for event in trigger.run()]
+
+ assert events == [
+ TriggerEvent({"status": expected_status, "job_name": JOB_NAME,
"message": expected_message})
+ ]
+
class TestSagemakerPipelineTrigger:
def test_serialize(self):
