This is an automated email from the ASF dual-hosted git repository.
o-nikolas pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 90f064fa244 Allow BedrockBatchInferenceOperator to wait until the job
is scheduled (#75230)
90f064fa244 is described below
commit 90f064fa244c8d81f29a92a955a7246b506f8f93
Author: Ramit Kataria <[email protected]>
AuthorDate: Wed Oct 21 10:07:16 2026 -0700
Allow BedrockBatchInferenceOperator to wait until the job is scheduled
(#75230)
BedrockBatchInferenceSensor can wait for a batch inference job to be either
scheduled or completed, but the operator could only wait for completion, in
both wait_for_completion and deferrable mode. Batch inference jobs sit in a
queue for an unpredictable amount of time, so many Dags only need to know
that
Bedrock has accepted the job, and today that takes a separate sensor task.
The
waiter and trigger for the scheduled state already exist.
---
providers/amazon/docs/operators/bedrock.rst | 3 ++
.../providers/amazon/aws/operators/bedrock.py | 27 +++++++++---
.../unit/amazon/aws/operators/test_bedrock.py | 48 ++++++++++++++++++++--
3 files changed, 68 insertions(+), 10 deletions(-)
diff --git a/providers/amazon/docs/operators/bedrock.rst
b/providers/amazon/docs/operators/bedrock.rst
index 74af2f116a7..5bb49aa891c 100644
--- a/providers/amazon/docs/operators/bedrock.rst
+++ b/providers/amazon/docs/operators/bedrock.rst
@@ -346,6 +346,9 @@ NOTE: Jobs are added to a queue and processed in order.
Given the potential wai
and the fact that the optional timeout parameter is measured in hours,
deferrable mode is
recommended over "wait_for_completion" in this case.
+By default the operator waits for the job to complete. Set
``success_state="scheduled"`` to stop
+waiting once Bedrock has queued the job. This applies to both
"wait_for_completion" and deferrable mode.
+
Example using an Amazon Bedrock Batch Inference Job:
.. exampleinclude::
/../../amazon/tests/system/amazon/aws/example_bedrock_batch_inference.py
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 3e0f0bc0b7a..fe376405d1e 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
@@ -36,6 +36,7 @@ from airflow.providers.amazon.aws.triggers.bedrock import (
BedrockAgentRuntimeDeletedTrigger,
BedrockAgentRuntimeReadyTrigger,
BedrockBatchInferenceCompletedTrigger,
+ BedrockBatchInferenceScheduledTrigger,
BedrockCustomizeModelCompletedTrigger,
BedrockIngestionJobTrigger,
BedrockKnowledgeBaseActiveTrigger,
@@ -47,6 +48,7 @@ from airflow.providers.common.compat.sdk import
AirflowException, conf, timezone
from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.triggers.bedrock import
BedrockBaseBatchInferenceTrigger
from airflow.sdk import Context
@@ -1239,12 +1241,15 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
: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)
+ :param wait_for_completion: Whether to wait for the job to reach
``success_state``. (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 success_state: The job state to wait for, either ``"completed"`` or
``"scheduled"``. Use
+ ``"scheduled"`` to stop waiting once Bedrock has queued the job.
Accepts the
+ ``BedrockBatchInferenceSensor.SuccessState`` values. (default:
"completed")
: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: 20)
- :param deferrable: If True, the operator will wait asynchronously for the
cluster to stop.
+ :param deferrable: If True, the operator will wait asynchronously for the
job to reach ``success_state``.
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.
@@ -1268,6 +1273,10 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
"output_uri",
"invoke_kwargs",
)
+ _WAITERS: dict[str, tuple[str, type[BedrockBaseBatchInferenceTrigger]]] = {
+ "completed": ("batch_inference_complete",
BedrockBatchInferenceCompletedTrigger),
+ "scheduled": ("batch_inference_scheduled",
BedrockBatchInferenceScheduledTrigger),
+ }
def __init__(
self,
@@ -1278,6 +1287,7 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
output_uri: str,
invoke_kwargs: dict[str, Any] | None = None,
wait_for_completion: bool = True,
+ success_state: str = "completed",
waiter_delay: int = 60,
waiter_max_attempts: int = 20,
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
@@ -1291,7 +1301,11 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
self.output_uri = output_uri
self.invoke_kwargs = invoke_kwargs or {}
+ if success_state not in self._WAITERS:
+ raise ValueError(f"success_state must be one of
{list(self._WAITERS)}, got {success_state!r}.")
+
self.wait_for_completion = wait_for_completion
+ self.success_state = success_state
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
self.deferrable = deferrable
@@ -1304,7 +1318,7 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
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"])
+ self.log.info("%s '%s' is %s.", self.activity,
validated_event["job_arn"], self.success_state)
return validated_event["job_arn"]
@@ -1320,11 +1334,12 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
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."
+ waiter_name, trigger_class = self._WAITERS[self.success_state]
+ task_description = f"for {self.activity} '{self.job_name}' to be
{self.success_state}."
if self.deferrable:
self.log.info("Deferring %s", task_description)
self.defer(
- trigger=BedrockBatchInferenceCompletedTrigger(
+ trigger=trigger_class(
job_arn=job_arn,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
@@ -1337,7 +1352,7 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
)
elif self.wait_for_completion:
self.log.info("Waiting %s", task_description)
- self.hook.get_waiter(waiter_name="batch_inference_complete").wait(
+ self.hook.get_waiter(waiter_name=waiter_name).wait(
jobIdentifier=job_arn,
WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts":
self.waiter_max_attempts},
)
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 cae3468a9d9..90468314a61 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
@@ -53,7 +53,11 @@ from airflow.providers.amazon.aws.operators.bedrock import (
BedrockRaGOperator,
BedrockUpdateGuardrailOperator,
)
-from airflow.providers.amazon.aws.triggers.bedrock import
BedrockAgentRuntimeDeletedTrigger
+from airflow.providers.amazon.aws.triggers.bedrock import (
+ BedrockAgentRuntimeDeletedTrigger,
+ BedrockBatchInferenceCompletedTrigger,
+ BedrockBatchInferenceScheduledTrigger,
+)
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
@@ -1093,8 +1097,8 @@ class TestBedrockBatchInferenceOperator:
hook = BedrockHook(aws_conn_id="aws_default")
yield hook
- def setup_method(self):
- self.operator = BedrockBatchInferenceOperator(
+ def _create_operator(self, **kwargs):
+ operator = BedrockBatchInferenceOperator(
task_id="test_task",
job_name=self.JOB_NAME,
role_arn=self.ROLE_ARN,
@@ -1102,8 +1106,13 @@ class TestBedrockBatchInferenceOperator:
input_uri=self.INPUT_URI,
output_uri=self.OUTPUT_URI,
invoke_kwargs=self.INVOKE_KWARGS,
+ **kwargs,
)
- self.operator.defer = mock.MagicMock()
+ operator.defer = mock.MagicMock()
+ return operator
+
+ def setup_method(self):
+ self.operator = self._create_operator()
@pytest.mark.parametrize(
("wait_for_completion", "deferrable"),
@@ -1126,6 +1135,37 @@ class TestBedrockBatchInferenceOperator:
assert bedrock_hook.get_waiter.call_count == wait_for_completion
assert self.operator.defer.call_count == deferrable
+ @pytest.mark.parametrize(
+ ("success_state", "waiter_name"),
+ [
+ pytest.param("completed", "batch_inference_complete",
id="completed"),
+ pytest.param("scheduled", "batch_inference_scheduled",
id="scheduled"),
+ ],
+ )
+ @mock.patch.object(BedrockHook, "get_waiter", autospec=True)
+ def test_wait_for_success_state(self, mock_get_waiter, success_state,
waiter_name, mock_conn):
+ self._create_operator(success_state=success_state).execute({})
+
+ assert mock_get_waiter.call_args.kwargs["waiter_name"] == waiter_name
+
+ @pytest.mark.parametrize(
+ ("success_state", "trigger_class"),
+ [
+ pytest.param("completed", BedrockBatchInferenceCompletedTrigger,
id="completed"),
+ pytest.param("scheduled", BedrockBatchInferenceScheduledTrigger,
id="scheduled"),
+ ],
+ )
+ def test_defer_for_success_state(self, success_state, trigger_class,
mock_conn):
+ operator = self._create_operator(success_state=success_state,
deferrable=True)
+
+ operator.execute({})
+
+ assert isinstance(operator.defer.call_args.kwargs["trigger"],
trigger_class)
+
+ def test_invalid_success_state(self):
+ with pytest.raises(ValueError, match="success_state must be one of"):
+ self._create_operator(success_state="running")
+
def test_template_fields(self):
validate_template_fields(self.operator)
