From 3c12d4594d0d853e046bf714bfccbb1a462e93c6 Mon Sep 17 00:00:00 2001 From: Hamza Tahir Date: Fri, 20 Dec 2024 22:37:35 +0100 Subject: [PATCH] Refactor error handling in SagemakerOrchestrator --- .../aws/orchestrators/sagemaker_orchestrator.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py b/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py index 29e0389494..367d2b060a 100644 --- a/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py +++ b/src/zenml/integrations/aws/orchestrators/sagemaker_orchestrator.py @@ -243,9 +243,6 @@ def prepare_or_run_pipeline( `boto3.Session` object. TypeError: If the network_config passed is not compatible with the AWS SageMaker NetworkConfig class. - KeyError: If required fields are missing from schedule_info. - ClientError: If there's an AWS API error. - BotoCoreError: If there's an error in the AWS SDK. Yields: A dictionary of metadata related to the pipeline run. @@ -567,16 +564,17 @@ def prepare_or_run_pipeline( logger.info(f"Created/Updated IAM policy: {policy_name}") except (ClientError, BotoCoreError) as e: - logger.error( + logger.warning( f"Failed to update IAM policy: {e}. " - f"Please ensure you have sufficient IAM permissions." + f"Please ensure your execution role has sufficient permissions " + f"to start pipeline executions." ) - raise except KeyError as e: - logger.error( - f"Missing required field for IAM policy creation: {e}" + logger.warning( + f"Missing required field for IAM policy creation: {e}. " + f"Please ensure your execution role has sufficient permissions " + f"to start pipeline executions." ) - raise # Create the EventBridge rule events_client = session.boto_session.client("events")