Add Bedrock Batch Inference Operator and accompanying parts

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
 ---------
 

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

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)