From 06c00dedd5c0ff7bf5dfe09a825c0f83cb43e096 Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 8 Apr 2025 10:34:21 -0700 Subject: [PATCH 1/2] Bedrock Batch Inference - Trying to stop a completed job is a successful result --- .../system/amazon/aws/example_bedrock_batch_inference.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py b/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py index ed0da3b0aea06..6ee5a934bd950 100644 --- a/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py +++ b/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py @@ -21,6 +21,8 @@ from datetime import datetime from tempfile import NamedTemporaryFile +from botocore.exceptions import ClientError + from airflow.decorators import task from airflow.models.baseoperator import chain from airflow.models.dag import DAG @@ -92,7 +94,12 @@ def generate_prompts(_env_id: str, _bucket: str, _key: str): @task(trigger_rule=TriggerRule.ALL_DONE) def stop_batch_inference(job_arn: str): log.info("Stopping Batch Inference Job.") - BedrockHook().conn.stop_model_invocation_job(jobIdentifier=job_arn) + try: + BedrockHook().conn.stop_model_invocation_job(jobIdentifier=job_arn) + except ClientError as e: + # If the job has already completed, boto will raise a ValidationException. Consider that a successful result. + if (e.response["Error"]["Code"] == "ValidationException") and ("State was: Completed" in e.response["Error"]["Message"]): + pass with DAG( From 4ba230158d5078f6dfcc63b298438bdc75aaccf7 Mon Sep 17 00:00:00 2001 From: ferruzzi Date: Tue, 8 Apr 2025 11:33:50 -0700 Subject: [PATCH 2/2] static checks --- .../system/amazon/aws/example_bedrock_batch_inference.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py b/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py index 6ee5a934bd950..2f1c18409e2b7 100644 --- a/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py +++ b/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py @@ -98,7 +98,9 @@ def stop_batch_inference(job_arn: str): BedrockHook().conn.stop_model_invocation_job(jobIdentifier=job_arn) except ClientError as e: # If the job has already completed, boto will raise a ValidationException. Consider that a successful result. - if (e.response["Error"]["Code"] == "ValidationException") and ("State was: Completed" in e.response["Error"]["Message"]): + if (e.response["Error"]["Code"] == "ValidationException") and ( + "State was: Completed" in e.response["Error"]["Message"] + ): pass