retry_policies.py

from sagemaker.workflow.retry import (
    StepRetryPolicy,
    StepExceptionTypeEnum,
    SageMakerJobStepRetryPolicy,
    SageMakerJobExceptionTypeEnum
)

step_retry_policy = StepRetryPolicy(
    exception_types=[
        StepExceptionTypeEnum.SERVICE_FAULT,
        StepExceptionTypeEnum.THROTTLING,
    ],
    backoff_rate=2.0,
    interval_seconds=30,
    expire_after_mins=240  # keep trying for for 4 hours max
)

job_retry_policy = SageMakerJobStepRetryPolicy(
    exception_types=[SageMakerJobExceptionTypeEnum.RESOURCE_LIMIT],
    failure_reason_types=[
        SageMakerJobExceptionTypeEnum.INTERNAL_ERROR,
        SageMakerJobExceptionTypeEnum.CAPACITY_ERROR,
    ],
    backoff_rate=2.0,
    interval_seconds=30,
    expire_after_mins=240  # keep trying for for 4 hours max
)