diff --git a/providers/amazon/docs/operators/bedrock.rst b/providers/amazon/docs/operators/bedrock.rst index 0cbd0a8cf120e..259b00b294b5b 100644 --- a/providers/amazon/docs/operators/bedrock.rst +++ b/providers/amazon/docs/operators/bedrock.rst @@ -74,6 +74,14 @@ To invoke a Claude V2 model using the Completions API you would use: :start-after: [START howto_operator_invoke_claude_model] :end-before: [END howto_operator_invoke_claude_model] +To invoke a Claude V3 Sonnet model using the Messages API you would use: + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_invoke_claude_messages] + :end-before: [END howto_operator_invoke_claude_messages] + .. _howto/operator:BedrockCustomizeModelOperator: @@ -237,6 +245,29 @@ Example using a PDF file in an Amazon S3 Bucket: :start-after: [START howto_operator_bedrock_external_sources_rag] :end-before: [END howto_operator_bedrock_external_sources_rag] +.. _howto/operator:BedrockBatchInferenceOperator: + +Create an Amazon Bedrock Batch Inference Job +============================================ + +To creates a batch inference job to invoke a model on multiple prompts, you can use +:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockBatchInferenceOperator`. + +The input must be formatted in jsonl and uploaded to an Amazon S3 bucket. Please see +https://docs.aws.amazon.com/bedrock/latest/userguide/batch-inference.html for details. + +NOTE: Jobs are added to a queue and processed in order. Given the potential wait times, +and the fact that the optional timeout parameter is measured in hours, deferrable mode is +recommended over "wait_for_completion" in this case. + +Example using an Amazon Bedrock Batch Inference Job: + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py + :language: python + :dedent: 4 + :start-after: [START howto_operator_bedrock_batch_inference] + :end-before: [END howto_operator_bedrock_batch_inference] + Sensors ------- @@ -298,6 +329,24 @@ To wait on the state of an Amazon Bedrock data ingestion job until it reaches a :start-after: [START howto_sensor_bedrock_ingest_data] :end-before: [END howto_sensor_bedrock_ingest_data] +.. _howto/sensor:BedrockBatchInferenceSensor: + +Wait for an Amazon Bedrock batch inference job +============================================== + +To wait on the state of an Amazon Bedrock batch inference job until it reaches the "Scheduled" or "Completed" +state you can use :class:`~airflow.providers.amazon.aws.sensors.bedrock.BedrockBatchInferenceScheduledSensor` + +Bedrock adds batch inference jobs to a queue, and they may take some time to complete. If you want to wait +for the job to complete, use TargetState.COMPLETED for the success_state, but if you only want to wait until +the service confirms that the job is in the queue, use TargetState.SCHEDULED. + +.. exampleinclude:: /../../amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py + :language: python + :dedent: 4 + :start-after: [START howto_sensor_bedrock_batch_inference_scheduled] + :end-before: [END howto_sensor_bedrock_batch_inference_scheduled] + Reference --------- diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py index ac8e86b92c807..f9715114b7408 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py @@ -33,6 +33,7 @@ ) from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator from airflow.providers.amazon.aws.triggers.bedrock import ( + BedrockBatchInferenceCompletedTrigger, BedrockCustomizeModelCompletedTrigger, BedrockIngestionJobTrigger, BedrockKnowledgeBaseActiveTrigger, @@ -869,3 +870,121 @@ def execute(self, context: Context) -> Any: self.log.info("\nQuery: %s\nRetrieved: %s", self.retrieval_query, result["retrievalResults"]) return result + + +class BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]): + """ + Create a batch inference job to invoke a model on multiple prompts. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:BedrockBatchInferenceOperator` + + :param job_name: A name to give the batch inference job. (templated) + :param role_arn: The ARN of the IAM role with permissions to create the knowledge base. (templated) + :param model_id: Name or ARN of the model to associate with this provisioned throughput. (templated) + :param input_uri: The S3 location of the input data. (templated) + :param output_uri: The S3 location of the output data. (templated) + :param invoke_kwargs: Additional keyword arguments to pass to the API call. (templated) + + :param wait_for_completion: Whether to wait for cluster to stop. (default: True) + NOTE: The way batch inference jobs work, your jobs are added to a queue and done "eventually" + so using deferrable mode is much more practical than using wait_for_completion. + :param waiter_delay: Time in seconds to wait between status checks. (default: 60) + :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 10) + :param deferrable: If True, the operator will wait asynchronously for the cluster to stop. + This implies waiting for completion. This mode requires aiobotocore module to be installed. + (default: False) + :param aws_conn_id: The Airflow connection used for AWS credentials. + If this is ``None`` or empty then the default boto3 behaviour is used. If + running Airflow in a distributed manner and aws_conn_id is None or + empty, then default boto3 configuration would be used (and must be + maintained on each worker node). + :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html + """ + + aws_hook_class = BedrockHook + template_fields: Sequence[str] = aws_template_fields( + "job_name", + "role_arn", + "model_id", + "input_uri", + "output_uri", + "invoke_kwargs", + ) + + def __init__( + self, + job_name: str, + role_arn: str, + model_id: str, + input_uri: str, + output_uri: str, + invoke_kwargs: dict[str, Any] | None = None, + wait_for_completion: bool = True, + waiter_delay: int = 60, + waiter_max_attempts: int = 10, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **kwargs, + ): + super().__init__(**kwargs) + self.job_name = job_name + self.role_arn = role_arn + self.model_id = model_id + self.input_uri = input_uri + self.output_uri = output_uri + self.invoke_kwargs = invoke_kwargs or {} + + self.wait_for_completion = wait_for_completion + self.waiter_delay = waiter_delay + self.waiter_max_attempts = waiter_max_attempts + self.deferrable = deferrable + + self.activity = "Bedrock batch inference job" + + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> str: + validated_event = validate_execute_complete_event(event) + + if validated_event["status"] != "success": + raise AirflowException(f"Error while running {self.activity}: {validated_event}") + + self.log.info("%s '%s' complete.", self.activity, validated_event["job_arn"]) + + return validated_event["job_arn"] + + def execute(self, context: Context) -> str: + response = self.hook.conn.create_model_invocation_job( + jobName=self.job_name, + roleArn=self.role_arn, + modelId=self.model_id, + inputDataConfig={"s3InputDataConfig": {"s3Uri": self.input_uri}}, + outputDataConfig={"s3OutputDataConfig": {"s3Uri": self.output_uri}}, + **self.invoke_kwargs, + ) + job_arn = response["jobArn"] + self.log.info("%s '%s' started with ARN: %s", self.activity, self.job_name, job_arn) + + task_description = f"for {self.activity} '{self.job_name}' to complete." + if self.deferrable: + self.log.info("Deferring %s", task_description) + self.defer( + trigger=BedrockBatchInferenceCompletedTrigger( + job_arn=job_arn, + waiter_delay=self.waiter_delay, + waiter_max_attempts=self.waiter_max_attempts, + aws_conn_id=self.aws_conn_id, + ), + method_name="execute_complete", + ) + elif self.wait_for_completion: + self.log.info("Waiting %s", task_description) + self.hook.get_waiter(waiter_name="batch_inference_complete").wait( + jobIdentifier=job_arn, + WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts}, + ) + + return job_arn diff --git a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py index 541702f4774d0..c74c3e772fb3c 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py @@ -26,6 +26,8 @@ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook from airflow.providers.amazon.aws.sensors.base_aws import AwsBaseSensor from airflow.providers.amazon.aws.triggers.bedrock import ( + BedrockBatchInferenceCompletedTrigger, + BedrockBatchInferenceScheduledTrigger, BedrockCustomizeModelCompletedTrigger, BedrockIngestionJobTrigger, BedrockKnowledgeBaseActiveTrigger, @@ -34,6 +36,7 @@ from airflow.providers.amazon.aws.utils.mixins import aws_template_fields if TYPE_CHECKING: + from airflow.providers.amazon.aws.triggers.bedrock import BedrockBaseBatchInferenceTrigger from airflow.utils.context import Context @@ -368,3 +371,110 @@ def execute(self, context: Context) -> Any: ) else: super().execute(context=context) + + +class BedrockBatchInferenceSensor(BedrockBaseSensor[BedrockHook]): + """ + Poll the batch inference job status until it reaches a terminal state; fails if creation fails. + + .. seealso:: + For more information on how to use this sensor, take a look at the guide: + :ref:`howto/sensor:BedrockBatchInferenceSensor` + + :param job_arn: The Amazon Resource Name (ARN) of the batch inference job. (templated) + :param success_state: A BedrockBatchInferenceSensor.TargetState; defaults to 'SCHEDULED' (templated) + + :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore + module to be installed. + (default: False, but can be overridden in config file by setting default_deferrable to True) + :param poke_interval: Polling period in seconds to check for the status of the job. (default: 5) + :param max_retries: Number of times before returning the current state (default: 24) + :param aws_conn_id: The Airflow connection used for AWS credentials. + If this is ``None`` or empty then the default boto3 behaviour is used. If + running Airflow in a distributed manner and aws_conn_id is None or + empty, then default boto3 configuration would be used (and must be + maintained on each worker node). + :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. + :param verify: Whether or not to verify SSL certificates. See: + https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html + :param botocore_config: Configuration dictionary (key-values) for botocore client. See: + https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html + """ + + class SuccessState: + """ + Target state for the BedrockBatchInferenceSensor. + + Bedrock adds batch inference jobs to a queue, and they may take some time to complete. + If you want to wait for the job to complete, use TargetState.COMPLETED, but if you only want + to wait until the service confirms that the job is in the queue, use TargetState.SCHEDULED. + + The normal successful progression of states is: + Submitted > Validating > Scheduled > InProgress > PartiallyCompleted > Completed + """ + + SCHEDULED = "scheduled" + COMPLETED = "completed" + + INTERMEDIATE_STATES: tuple[str, ...] # Defined in __init__ based on target state + FAILURE_STATES: tuple[str, ...] = ("Failed", "Stopped", "PartiallyCompleted", "Expired") + SUCCESS_STATES: tuple[str, ...] # Defined in __init__ based on target state + FAILURE_MESSAGE = "Bedrock batch inference job sensor failed." + INVALID_SUCCESS_STATE_MESSAGE = "success_state must be an instance of TargetState." + + aws_hook_class = BedrockHook + + template_fields: Sequence[str] = aws_template_fields("job_arn", "success_state") + + def __init__( + self, + *, + job_arn: str, + success_state: SuccessState | str = SuccessState.SCHEDULED, + poke_interval: int = 120, + max_retries: int = 75, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.poke_interval = poke_interval + self.max_retries = max_retries + self.job_arn = job_arn + self.success_state = success_state + + base_success_states: tuple[str, ...] = ("Completed",) + base_intermediate_states: tuple[str, ...] = ("Submitted", "InProgress", "Stopping", "Validating") + scheduled_state = ("Scheduled",) + self.trigger_class: type[BedrockBaseBatchInferenceTrigger] + + if self.success_state == BedrockBatchInferenceSensor.SuccessState.COMPLETED: + intermediate_states = base_intermediate_states + scheduled_state + success_states = base_success_states + self.trigger_class = BedrockBatchInferenceCompletedTrigger + elif self.success_state == BedrockBatchInferenceSensor.SuccessState.SCHEDULED: + intermediate_states = base_intermediate_states + success_states = base_success_states + scheduled_state + self.trigger_class = BedrockBatchInferenceScheduledTrigger + else: + raise ValueError( + "Success states for BedrockBatchInferenceSensor must be set using a BedrockBatchInferenceSensor.SuccessState" + ) + + BedrockBatchInferenceSensor.INTERMEDIATE_STATES = intermediate_states or base_intermediate_states + BedrockBatchInferenceSensor.SUCCESS_STATES = success_states or base_success_states + + def get_state(self) -> str: + return self.hook.conn.get_model_invocation_job(jobIdentifier=self.job_arn)["status"] + + def execute(self, context: Context) -> Any: + if self.deferrable: + self.defer( + trigger=self.trigger_class( + job_arn=self.job_arn, + waiter_delay=int(self.poke_interval), + waiter_max_attempts=self.max_retries, + aws_conn_id=self.aws_conn_id, + ), + method_name="poke", + ) + else: + super().execute(context=context) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py b/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py index 99d632d26e0c3..faac0f90753e6 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py +++ b/providers/amazon/src/airflow/providers/amazon/aws/triggers/bedrock.py @@ -20,6 +20,7 @@ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger +from airflow.utils.types import NOTSET, ArgNotSet if TYPE_CHECKING: from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook @@ -182,3 +183,100 @@ def __init__( def hook(self) -> AwsGenericHook: return BedrockAgentHook(aws_conn_id=self.aws_conn_id) + + +class BedrockBaseBatchInferenceTrigger(AwsBaseWaiterTrigger): + """ + Trigger when a batch inference job is complete. + + :param job_arn: The Amazon Resource Name (ARN) of the batch inference job. + + :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120) + :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75) + :param aws_conn_id: The Airflow connection used for AWS credentials. + """ + + def __init__( + self, + *, + job_arn: str, + waiter_name: str | ArgNotSet = NOTSET, # This must be defined in the child class. + waiter_delay: int = 120, + waiter_max_attempts: int = 75, + aws_conn_id: str | None = None, + ) -> None: + if waiter_name == NOTSET: + raise NotImplementedError("Triggers must provide a waiter name.") + + super().__init__( + serialized_fields={"job_arn": job_arn}, + waiter_name=str(waiter_name), # Cast a string to a string to make mypy happy + waiter_args={"jobIdentifier": job_arn}, + failure_message="Bedrock batch inference job failed.", + status_message="Status of Bedrock batch inference job is", + status_queries=["status"], + return_key="job_arn", + return_value=job_arn, + waiter_delay=waiter_delay, + waiter_max_attempts=waiter_max_attempts, + aws_conn_id=aws_conn_id, + ) + + def hook(self) -> AwsGenericHook: + return BedrockHook(aws_conn_id=self.aws_conn_id) + + +class BedrockBatchInferenceCompletedTrigger(BedrockBaseBatchInferenceTrigger): + """ + Trigger when a batch inference job is complete. + + :param job_arn: The Amazon Resource Name (ARN) of the batch inference job. + + :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120) + :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75) + :param aws_conn_id: The Airflow connection used for AWS credentials. + """ + + def __init__( + self, + *, + job_arn: str, + waiter_delay: int = 120, + waiter_max_attempts: int = 75, + aws_conn_id: str | None = None, + ) -> None: + super().__init__( + waiter_name="batch_inference_complete", + job_arn=job_arn, + waiter_delay=waiter_delay, + waiter_max_attempts=waiter_max_attempts, + aws_conn_id=aws_conn_id, + ) + + +class BedrockBatchInferenceScheduledTrigger(BedrockBaseBatchInferenceTrigger): + """ + Trigger when a batch inference job is scheduled. + + :param job_arn: The Amazon Resource Name (ARN) of the batch inference job. + + :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120) + :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75) + :param aws_conn_id: The Airflow connection used for AWS credentials. + """ + + def __init__( + self, + *, + job_arn: str, + waiter_delay: int = 120, + waiter_max_attempts: int = 75, + aws_conn_id: str | None = None, + ) -> None: + super().__init__( + waiter_name="batch_inference_scheduled", + job_arn=job_arn, + waiter_delay=waiter_delay, + waiter_max_attempts=waiter_max_attempts, + aws_conn_id=aws_conn_id, + ) diff --git a/providers/amazon/src/airflow/providers/amazon/aws/waiters/bedrock.json b/providers/amazon/src/airflow/providers/amazon/aws/waiters/bedrock.json index c913b4dc7c68a..18a6294a12810 100644 --- a/providers/amazon/src/airflow/providers/amazon/aws/waiters/bedrock.json +++ b/providers/amazon/src/airflow/providers/amazon/aws/waiters/bedrock.json @@ -68,6 +68,140 @@ "state": "failure" } ] + }, + "batch_inference_complete": { + "delay": 120, + "maxAttempts": 75, + "operation": "GetModelInvocationJob", + "acceptors": [ + { + "matcher": "path", + "argument": "status", + "expected": "Completed", + "state": "success" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Failed", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Stopped", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "PartiallyCompleted", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Expired", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Stopping", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Submitted", + "state": "retry" + }, + { + "matcher": "path", + "argument": "status", + "expected": "InProgress", + "state": "retry" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Validating", + "state": "retry" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Scheduled", + "state": "retry" + } + ] + }, + "batch_inference_scheduled": { + "delay": 120, + "maxAttempts": 75, + "operation": "GetModelInvocationJob", + "acceptors": [ + { + "matcher": "path", + "argument": "status", + "expected": "Completed", + "state": "success" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Failed", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Stopped", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Stopping", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "PartiallyCompleted", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Expired", + "state": "failure" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Submitted", + "state": "retry" + }, + { + "matcher": "path", + "argument": "status", + "expected": "InProgress", + "state": "retry" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Validating", + "state": "retry" + }, + { + "matcher": "path", + "argument": "status", + "expected": "Scheduled", + "state": "success" + } + ] } } } 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 new file mode 100644 index 0000000000000..ed0da3b0aea06 --- /dev/null +++ b/providers/amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py @@ -0,0 +1,182 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import json +import logging +from datetime import datetime +from tempfile import NamedTemporaryFile + +from airflow.decorators import task +from airflow.models.baseoperator import chain +from airflow.models.dag import DAG +from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook +from airflow.providers.amazon.aws.hooks.s3 import S3Hook +from airflow.providers.amazon.aws.operators.bedrock import ( + BedrockBatchInferenceOperator, + BedrockInvokeModelOperator, +) +from airflow.providers.amazon.aws.operators.s3 import ( + S3CreateBucketOperator, + S3DeleteBucketOperator, +) +from airflow.providers.amazon.aws.sensors.bedrock import BedrockBatchInferenceSensor +from airflow.utils.trigger_rule import TriggerRule + +from system.amazon.aws.utils import SystemTestContextBuilder + +log = logging.getLogger(__name__) + + +# Externally fetched variables: +ROLE_ARN_KEY = "ROLE_ARN" +sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build() + +DAG_ID = "example_bedrock_batch_inference" + +####################################################################### +# NOTE: +# Access to the following foundation model must be requested via +# the Amazon Bedrock console and may take up to 24 hours to apply: +####################################################################### + +CLAUDE_MODEL_ID = "anthropic.claude-3-sonnet-20240229-v1:0" +ANTHROPIC_VERSION = "bedrock-2023-05-31" + +# Batch inferences currently require a minimum of 100 prompts per batch. +MIN_NUM_PROMPTS = 300 +PROMPT_TEMPLATE = "Even numbers are red. Odd numbers are blue. What color is {n}?" + + +@task +def generate_prompts(_env_id: str, _bucket: str, _key: str): + """ + Bedrock Batch Inference requires one or more jsonl-formatted files in an S3 bucket. + + The JSONL format requires one serialized json object per prompt per line. + """ + with NamedTemporaryFile(mode="w") as tmp_file: + # Generate the required number of prompts. + prompts = [ + { + "modelInput": { + "anthropic_version": ANTHROPIC_VERSION, + "max_tokens": 1000, + "messages": [PROMPT_TEMPLATE.format(n=n)], + }, + } + for n in range(MIN_NUM_PROMPTS) + ] + + # Convert each prompt to serialized json, append a newline, and write that line to the temp file. + tmp_file.writelines(json.dumps(prompt) + "\n" for prompt in prompts) + + # Upload the file to S3. + S3Hook().conn.upload_file(tmp_file.name, _bucket, _key) + + +@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) + + +with DAG( + dag_id=DAG_ID, + schedule="@once", + start_date=datetime(2021, 1, 1), + tags={"example"}, + catchup=False, +) as dag: + test_context = sys_test_context_task() + env_id = test_context["ENV_ID"] + + bucket_name = f"{env_id}-bedrock" + input_data_s3_key = f"{env_id}/prompt_list.jsonl" + input_uri = f"s3://{bucket_name}/{input_data_s3_key}" + output_uri = f"s3://{bucket_name}/output/" + job_name = f"batch-infer-{env_id}" + + # Test that this configuration works for a single prompt before trying the batch inferences. + # [START howto_operator_invoke_claude_messages] + invoke_claude_messages = BedrockInvokeModelOperator( + task_id="invoke_claude_messages", + model_id=CLAUDE_MODEL_ID, + input_data={ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 1000, + "messages": [{"role": "user", "content": PROMPT_TEMPLATE.format(n=42)}], + }, + ) + # [END howto_operator_invoke_claude_messages] + + create_bucket = S3CreateBucketOperator(task_id="create_bucket", bucket_name=bucket_name) + + # [START howto_operator_bedrock_batch_inference] + batch_infer = BedrockBatchInferenceOperator( + task_id="batch_infer", + job_name=job_name, + role_arn=test_context[ROLE_ARN_KEY], + model_id=CLAUDE_MODEL_ID, + input_uri=input_uri, + output_uri=output_uri, + ) + # [END howto_operator_bedrock_batch_inference] + batch_infer.wait_for_completion = False + + # [START howto_sensor_bedrock_batch_inference_scheduled] + await_job_scheduled = BedrockBatchInferenceSensor( + task_id="await_job_scheduled", + job_arn=batch_infer.output, + success_state=BedrockBatchInferenceSensor.SuccessState.SCHEDULED, + ) + # [END howto_sensor_bedrock_batch_inference_scheduled] + + stop_job = stop_batch_inference(batch_infer.output) + + delete_bucket = S3DeleteBucketOperator( + task_id="delete_bucket", + trigger_rule=TriggerRule.ALL_DONE, + bucket_name=bucket_name, + force_delete=True, + ) + + chain( + # TEST SETUP + test_context, + invoke_claude_messages, + create_bucket, + generate_prompts(env_id, bucket_name, input_data_s3_key), + # TEST BODY + batch_infer, + await_job_scheduled, + stop_job, + # TEST TEARDOWN + delete_bucket, + ) + + from tests_common.test_utils.watcher import watcher + + # This test needs watcher in order to properly mark success/failure + # when "tearDown" task with trigger rule is part of the DAG + list(dag.tasks) >> watcher() + + +from tests_common.test_utils.system_tests import get_test_run # noqa: E402 + +# Needed to run the example DAG with pytest (see: tests/system/README.md#run_via_pytest) +test_run = get_test_run(dag) diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py index 527f0385e60df..4591ab97c1bdd 100644 --- a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py +++ b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py @@ -28,6 +28,7 @@ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook, BedrockRuntimeHook from airflow.providers.amazon.aws.operators.bedrock import ( + BedrockBatchInferenceOperator, BedrockCreateDataSourceOperator, BedrockCreateKnowledgeBaseOperator, BedrockCreateProvisionedModelThroughputOperator, @@ -549,3 +550,62 @@ def test_template_fields(self): vector_search_config=self.VECTOR_SEARCH_CONFIG, ) validate_template_fields(op) + + +class TestBedrockBatchInferenceOperator: + JOB_NAME = "job_name" + ROLE_ARN = "role_arn" + MODEL_ID = "model_id" + INPUT_URI = "input_uri" + OUTPUT_URI = "output_uri" + INVOKE_KWARGS = {"tags": {"key": "key", "value": "value"}} + + JOB_ARN = "job_arn" + + @pytest.fixture + def mock_conn(self) -> Generator[BaseAwsConnection, None, None]: + with mock.patch.object(BedrockHook, "conn") as _conn: + _conn.create_model_invocation_job.return_value = {"jobArn": self.JOB_ARN} + yield _conn + + @pytest.fixture + def bedrock_hook(self) -> Generator[BedrockHook, None, None]: + with mock_aws(): + hook = BedrockHook(aws_conn_id="aws_default") + yield hook + + def setup_method(self): + self.operator = BedrockBatchInferenceOperator( + task_id="test_task", + job_name=self.JOB_NAME, + role_arn=self.ROLE_ARN, + model_id=self.MODEL_ID, + input_uri=self.INPUT_URI, + output_uri=self.OUTPUT_URI, + invoke_kwargs=self.INVOKE_KWARGS, + ) + self.operator.defer = mock.MagicMock() + + @pytest.mark.parametrize( + "wait_for_completion, deferrable", + [ + pytest.param(False, False, id="no_wait"), + pytest.param(True, False, id="wait"), + pytest.param(False, True, id="defer"), + ], + ) + @mock.patch.object(BedrockHook, "get_waiter") + def test_customize_model_wait_combinations( + self, _, wait_for_completion, deferrable, mock_conn, bedrock_hook + ): + self.operator.wait_for_completion = wait_for_completion + self.operator.deferrable = deferrable + + response = self.operator.execute({}) + + assert response == self.JOB_ARN + assert bedrock_hook.get_waiter.call_count == wait_for_completion + assert self.operator.defer.call_count == deferrable + + def test_template_fields(self): + validate_template_fields(self.operator) diff --git a/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py b/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py index 52151cd953942..dc5ce64c4ad66 100644 --- a/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py +++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py @@ -21,9 +21,10 @@ import pytest -from airflow.exceptions import AirflowException +from airflow.exceptions import AirflowException, TaskDeferred from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook from airflow.providers.amazon.aws.sensors.bedrock import ( + BedrockBatchInferenceSensor, BedrockCustomizeModelCompletedSensor, BedrockIngestionJobSensor, BedrockKnowledgeBaseActiveSensor, @@ -244,3 +245,95 @@ def test_poke_failure_states(self, mock_conn, state): sensor = self.SENSOR(**self.default_op_kwargs, aws_conn_id=None) with pytest.raises(AirflowException, match=sensor.FAILURE_MESSAGE): sensor.poke({}) + + +class TestBedrockBatchInferenceSensor: + SENSOR = BedrockBatchInferenceSensor + + @pytest.fixture( + params=[ + BedrockBatchInferenceSensor.SuccessState.COMPLETED, + BedrockBatchInferenceSensor.SuccessState.SCHEDULED, + ] + ) + def success_state(self, request): + return request.param + + @pytest.fixture(params=["deferrable", "not deferrable"]) + def is_deferrable(self, request): + # I did it this way instead of passing True/False purely so the pytest names are more descriptive. + return request.param == "deferrable" + + def setup_method(self, is_deferrable): + self.default_op_kwargs = dict( + task_id="test_bedrock_batch_inference_sensor", + job_arn="job_arn", + deferrable=is_deferrable, + ) + + def test_base_aws_op_attributes(self, success_state): + op = self.SENSOR(**self.default_op_kwargs, success_state=success_state) + + if success_state == BedrockBatchInferenceSensor.SuccessState.COMPLETED: + assert "Scheduled" in op.INTERMEDIATE_STATES + assert "Scheduled" not in op.SUCCESS_STATES + elif success_state == BedrockBatchInferenceSensor.SuccessState.SCHEDULED: + assert "Scheduled" in op.SUCCESS_STATES + assert "Scheduled" not in op.INTERMEDIATE_STATES + assert op.hook.aws_conn_id == "aws_default" + assert op.hook._region_name is None + assert op.hook._verify is None + assert op.hook._config is None + + op = self.SENSOR( + **self.default_op_kwargs, + aws_conn_id="aws-test-custom-conn", + region_name="eu-west-1", + verify=False, + botocore_config={"read_timeout": 42}, + ) + assert op.hook.aws_conn_id == "aws-test-custom-conn" + assert op.hook._region_name == "eu-west-1" + assert op.hook._verify is False + assert op.hook._config is not None + assert op.hook._config.read_timeout == 42 + + @mock.patch.object(BedrockHook, "conn") + def test_poke_success_states(self, mock_conn, success_state, is_deferrable): + op = self.SENSOR(**self.default_op_kwargs, success_state=success_state) + + for state in op.SUCCESS_STATES: + mock_conn.get_model_invocation_job.return_value = {"status": state} + + if is_deferrable: + with pytest.raises(TaskDeferred): + op.execute({}) + else: + assert op.poke({}) is True + + @mock.patch.object(BedrockHook, "conn") + def test_poke_intermediate_states(self, mock_conn, success_state, is_deferrable): + op = self.SENSOR(**self.default_op_kwargs, success_state=success_state) + + for state in op.INTERMEDIATE_STATES: + mock_conn.get_model_invocation_job.return_value = {"status": state} + + if is_deferrable: + with pytest.raises(TaskDeferred): + op.execute({}) + else: + assert op.poke({}) is False + + @mock.patch.object(BedrockHook, "conn") + def test_poke_failure_states(self, mock_conn, success_state, is_deferrable): + op = self.SENSOR(**self.default_op_kwargs, success_state=success_state) + + for state in op.FAILURE_STATES: + mock_conn.get_model_invocation_job.return_value = {"status": state} + + if is_deferrable: + with pytest.raises(AirflowException, match=op.FAILURE_MESSAGE): + op.poke({}) + else: + with pytest.raises(TaskDeferred): + op.execute({}) diff --git a/providers/amazon/tests/unit/amazon/aws/triggers/test_bedrock.py b/providers/amazon/tests/unit/amazon/aws/triggers/test_bedrock.py index be17bb958d0f4..963f96211e757 100644 --- a/providers/amazon/tests/unit/amazon/aws/triggers/test_bedrock.py +++ b/providers/amazon/tests/unit/amazon/aws/triggers/test_bedrock.py @@ -23,6 +23,8 @@ from airflow.providers.amazon.aws.hooks.bedrock import BedrockAgentHook, BedrockHook from airflow.providers.amazon.aws.triggers.bedrock import ( + BedrockBatchInferenceCompletedTrigger, + BedrockBatchInferenceScheduledTrigger, BedrockCustomizeModelCompletedTrigger, BedrockIngestionJobTrigger, BedrockKnowledgeBaseActiveTrigger, @@ -169,3 +171,63 @@ async def test_run_success(self, mock_async_conn, mock_get_waiter): assert_expected_waiter_type(mock_get_waiter, self.EXPECTED_WAITER_NAME) assert response == TriggerEvent({"status": "success", "ingestion_job_id": self.INGESTION_JOB_ID}) mock_get_waiter().wait.assert_called_once() + + +class TestBedrockBatchInferenceCompletedTrigger(TestBaseBedrockTrigger): + EXPECTED_WAITER_NAME = "batch_inference_complete" + + JOB_ARN = "job_arn" + + def test_serialization(self): + """Assert that arguments and classpath are correctly serialized.""" + trigger = BedrockBatchInferenceCompletedTrigger(job_arn=self.JOB_ARN) + + classpath, kwargs = trigger.serialize() + + assert classpath == BASE_TRIGGER_CLASSPATH + "BedrockBatchInferenceCompletedTrigger" + assert kwargs.get("job_arn") == self.JOB_ARN + + @pytest.mark.asyncio + @mock.patch.object(BedrockHook, "get_waiter") + @mock.patch.object(BedrockHook, "get_async_conn") + async def test_run_success(self, mock_async_conn, mock_get_waiter): + mock_async_conn.__aenter__.return_value = mock.MagicMock() + mock_get_waiter().wait = AsyncMock() + trigger = BedrockBatchInferenceCompletedTrigger(job_arn=self.JOB_ARN) + + generator = trigger.run() + response = await generator.asend(None) + + assert_expected_waiter_type(mock_get_waiter, self.EXPECTED_WAITER_NAME) + assert response == TriggerEvent({"status": "success", "job_arn": self.JOB_ARN}) + mock_get_waiter().wait.assert_called_once() + + +class TestBedrockBatchInferenceScheduledTrigger(TestBaseBedrockTrigger): + EXPECTED_WAITER_NAME = "batch_inference_scheduled" + + JOB_ARN = "job_arn" + + def test_serialization(self): + """Assert that arguments and classpath are correctly serialized.""" + trigger = BedrockBatchInferenceScheduledTrigger(job_arn=self.JOB_ARN) + + classpath, kwargs = trigger.serialize() + + assert classpath == BASE_TRIGGER_CLASSPATH + "BedrockBatchInferenceScheduledTrigger" + assert kwargs.get("job_arn") == self.JOB_ARN + + @pytest.mark.asyncio + @mock.patch.object(BedrockHook, "get_waiter") + @mock.patch.object(BedrockHook, "get_async_conn") + async def test_run_success(self, mock_async_conn, mock_get_waiter): + mock_async_conn.__aenter__.return_value = mock.MagicMock() + mock_get_waiter().wait = AsyncMock() + trigger = BedrockBatchInferenceScheduledTrigger(job_arn=self.JOB_ARN) + + generator = trigger.run() + response = await generator.asend(None) + + assert_expected_waiter_type(mock_get_waiter, self.EXPECTED_WAITER_NAME) + assert response == TriggerEvent({"status": "success", "job_arn": self.JOB_ARN}) + mock_get_waiter().wait.assert_called_once() diff --git a/providers/amazon/tests/unit/amazon/aws/waiters/test_bedrock.py b/providers/amazon/tests/unit/amazon/aws/waiters/test_bedrock.py index 3d8a3a1af1f0f..3214e662493b2 100644 --- a/providers/amazon/tests/unit/amazon/aws/waiters/test_bedrock.py +++ b/providers/amazon/tests/unit/amazon/aws/waiters/test_bedrock.py @@ -17,6 +17,9 @@ from __future__ import annotations +import inspect +import re +import sys from unittest import mock import boto3 @@ -25,6 +28,7 @@ from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook from airflow.providers.amazon.aws.sensors.bedrock import ( + BedrockBatchInferenceSensor, BedrockCustomizeModelCompletedSensor, BedrockProvisionModelThroughputCompletedSensor, ) @@ -32,8 +36,17 @@ class TestBedrockCustomWaiters: def test_service_waiters(self): - assert "model_customization_job_complete" in BedrockHook().list_waiters() - assert "provisioned_model_throughput_complete" in BedrockHook().list_waiters() + """Ensure that all custom Bedrock waiters have unit tests.""" + + def _class_tests_a_waiter(class_name: str) -> bool: + """Check if the class name starts with 'Test' and ends with 'Waiter'.""" + return bool(re.match(r"^Test[A-Za-z]+Waiter$", class_name)) + + # Collect WAITER_NAME from each waiter test class in this module. + test_classes = inspect.getmembers(sys.modules[__name__], inspect.isclass) + waiters_tested = [cls.WAITER_NAME for (name, cls) in test_classes if _class_tests_a_waiter(name)] + + assert sorted(BedrockHook()._list_custom_waiters()) == sorted(waiters_tested) class TestBedrockCustomWaitersBase: @@ -103,3 +116,75 @@ def test_model_customization_job_wait(self, mock_get_job): BedrockHook().get_waiter(self.WAITER_NAME).wait( jobIdentifier="job_id", WaiterConfig={"Delay": 0.01, "MaxAttempts": 3} ) + + +class TestBatchInferenceCompleteWaiter(TestBedrockCustomWaitersBase): + WAITER_NAME = "batch_inference_complete" + SENSOR = BedrockBatchInferenceSensor( + task_id="task_id", + job_arn="job_arn", + success_state=BedrockBatchInferenceSensor.SuccessState.COMPLETED, + ) + + @pytest.fixture + def mock_get_job(self): + with mock.patch.object(self.client, "get_model_invocation_job") as mock_getter: + yield mock_getter + + @pytest.mark.parametrize("state", SENSOR.SUCCESS_STATES) + def test_batch_inference_complete(self, state, mock_get_job): + mock_get_job.return_value = {"status": state} + + BedrockHook().get_waiter(self.WAITER_NAME).wait(jobIdentifier="job_arn") + + @pytest.mark.parametrize("state", SENSOR.FAILURE_STATES) + def test_batch_inference_failed(self, state, mock_get_job): + mock_get_job.return_value = {"status": state} + + with pytest.raises(botocore.exceptions.WaiterError): + BedrockHook().get_waiter(self.WAITER_NAME).wait(jobIdentifier="job_arn") + + def test_batch_inference_wait(self, mock_get_job): + wait = {"status": "InProgress"} + success = {"status": "Completed"} + mock_get_job.side_effect = [wait, wait, success] + + BedrockHook().get_waiter(self.WAITER_NAME).wait( + jobIdentifier="job_arn", WaiterConfig={"Delay": 0.01, "MaxAttempts": 3} + ) + + +class TestBatchInferenceScheduledWaiter(TestBedrockCustomWaitersBase): + WAITER_NAME = "batch_inference_scheduled" + SENSOR = BedrockBatchInferenceSensor( + task_id="task_id", + job_arn="job_arn", + success_state=BedrockBatchInferenceSensor.SuccessState.SCHEDULED, + ) + + @pytest.fixture + def mock_get_job(self): + with mock.patch.object(self.client, "get_model_invocation_job") as mock_getter: + yield mock_getter + + @pytest.mark.parametrize("state", SENSOR.SUCCESS_STATES) + def test_batch_inference_complete(self, state, mock_get_job): + mock_get_job.return_value = {"status": state} + + BedrockHook().get_waiter(self.WAITER_NAME).wait(jobIdentifier="job_arn") + + @pytest.mark.parametrize("state", SENSOR.FAILURE_STATES) + def test_batch_inference_failed(self, state, mock_get_job): + mock_get_job.return_value = {"status": state} + + with pytest.raises(botocore.exceptions.WaiterError): + BedrockHook().get_waiter(self.WAITER_NAME).wait(jobIdentifier="job_arn") + + def test_batch_inference_wait(self, mock_get_job): + wait = {"status": "InProgress"} + success = {"status": "Completed"} + mock_get_job.side_effect = [wait, wait, success] + + BedrockHook().get_waiter(self.WAITER_NAME).wait( + jobIdentifier="job_arn", WaiterConfig={"Delay": 0.01, "MaxAttempts": 3} + )