This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 90efd88e2c1 Fix BedrockBatchInferenceSensor sharing success states
across tasks (#75265)
90efd88e2c1 is described below
commit 90efd88e2c1fe5ce22e8bf83d4051af32b62d509
Author: Ramit Kataria <[email protected]>
AuthorDate: Fri Oct 23 11:49:40 2026 -0700
Fix BedrockBatchInferenceSensor sharing success states across tasks (#75265)
The sensor stored the states for its success_state on the class, so every
BedrockBatchInferenceSensor in a process used the states of whichever one
was
created last. Parsing a Dag file creates all of its tasks, so a sensor
waiting
for completion could succeed as soon as the job was scheduled if the same
file
also had a sensor waiting for the scheduled state, and the other way around.
---
.../src/airflow/providers/amazon/aws/sensors/bedrock.py | 4 ++--
.../amazon/tests/unit/amazon/aws/sensors/test_bedrock.py | 12 ++++++++++++
2 files changed, 14 insertions(+), 2 deletions(-)
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 c8824e1f574..d0a38015463 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py
@@ -492,8 +492,8 @@ class
BedrockBatchInferenceSensor(BedrockBaseSensor[BedrockHook]):
"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
+ self.INTERMEDIATE_STATES = intermediate_states
+ self.SUCCESS_STATES = success_states
def get_state(self) -> str:
return
self.hook.conn.get_model_invocation_job(jobIdentifier=self.job_arn)["status"]
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 fd8da50f58d..1a8c37186a0 100644
--- a/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py
@@ -358,3 +358,15 @@ class TestBedrockBatchInferenceSensor:
else:
with pytest.raises(TaskDeferred):
op.execute({})
+
+ @mock.patch.object(BedrockHook, "conn")
+ def test_success_state_not_shared_between_instances(self, mock_conn):
+ op = self.SENSOR(
+ **self.default_op_kwargs,
success_state=BedrockBatchInferenceSensor.SuccessState.COMPLETED
+ )
+ self.SENSOR(
+ **self.default_op_kwargs,
success_state=BedrockBatchInferenceSensor.SuccessState.SCHEDULED
+ )
+ mock_conn.get_model_invocation_job.return_value = {"status":
"Scheduled"}
+
+ assert op.poke({}) is False
