SameerMesiah97 commented on code in PR #76848:
URL: https://github.com/apache/airflow/pull/76848#discussion_r4529634260
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/bedrock.py:
##########
@@ -83,6 +85,12 @@ def poke(self, context: Context, **kwargs) -> bool:
return state not in self.INTERMEDIATE_STATES
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
Review Comment:
Could we verify that the waiter and `poke()` agree on which resource states
mean success, which mean failure, and which should keep polling? For Bedrock
batch inference, this should cover both the `scheduled` and `completed` targets.
##########
providers/amazon/tests/unit/amazon/aws/sensors/test_bedrock.py:
##########
@@ -370,3 +371,47 @@ def test_success_state_not_shared_between_instances(self,
mock_conn):
mock_conn.get_model_invocation_job.return_value = {"status":
"Scheduled"}
assert op.poke({}) is False
+
+
+SENSORS = [
+ pytest.param(BedrockCustomizeModelCompletedSensor, {"job_name":
"job_name"}, id="customize_model"),
+ pytest.param(
+ BedrockProvisionModelThroughputCompletedSensor, {"model_id":
"model_id"}, id="provision_throughput"
+ ),
+ pytest.param(
+ BedrockKnowledgeBaseActiveSensor, {"knowledge_base_id":
"knowledge_base_id"}, id="knowledge_base"
+ ),
+ pytest.param(
+ BedrockIngestionJobSensor,
+ {
+ "knowledge_base_id": "knowledge_base_id",
+ "data_source_id": "data_source_id",
+ "ingestion_job_id": "ingestion_job_id",
+ },
+ id="ingestion_job",
+ ),
+ pytest.param(BedrockBatchInferenceSensor, {"job_arn": "job_arn"},
id="batch_inference"),
+]
+
+
+class TestBedrockSensorDeferral:
+ @pytest.mark.parametrize(("sensor_class", "kwargs"), SENSORS)
+ def test_execute_defers_to_execute_complete(self, sensor_class, kwargs):
+ sensor = sensor_class(task_id="test_bedrock_sensor", deferrable=True,
**kwargs)
+
+ with pytest.raises(TaskDeferred) as deferred:
+ sensor.execute({})
+
+ assert deferred.value.method_name == "execute_complete"
+
+ def test_execute_complete_success_event(self):
+ sensor = BedrockBatchInferenceSensor(task_id="test_bedrock_sensor",
job_arn="job_arn")
+
+ assert sensor.execute_complete({}, event={"status": "success",
"job_arn": "job_arn"}) is None
+
+ def test_execute_complete_error_event(self):
+ sensor = BedrockBatchInferenceSensor(task_id="test_bedrock_sensor",
job_arn="job_arn")
+ event = {"status": "error", "message": "Waiter error: max attempts
reached", "job_arn": "job_arn"}
+
+ with pytest.raises(SensorTriggerFailedError, match="max attempts
reached"):
+ sensor.execute_complete({}, event=event)
Review Comment:
Could we add (or extend) a test to cover the specific scenario you mentioned
in the PR description? You could keep the mocked Bedrock batch job in
Submitted, let the trigger exhaust its attempts, then pass its event to the
resume method recorded in `TaskDeferred`. Assert that resuming raises
`SensorTriggerFailedError`
##########
providers/amazon/src/airflow/providers/amazon/aws/sensors/comprehend.py:
##########
@@ -227,11 +235,23 @@ def execute(self, context: Context) -> Any:
verify=self.verify,
botocore_config=self.botocore_config,
),
- method_name="poke",
+ method_name="execute_complete",
)
else:
super().execute(context=context)
+ def execute_complete(self, context: Context, event: dict[str, Any] | None
= None) -> None:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event["status"] != "success":
Review Comment:
The state checks here could be more granular. There are more possible states
besides 'success' and not 'success'.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
