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 1159db91258 Deprecate deferrable implying wait in no-wait-default AWS
operators (#75361)
1159db91258 is described below
commit 1159db91258a045d4687a4d9413b4c9c906144f0
Author: Ramit Kataria <[email protected]>
AuthorDate: Fri Oct 23 11:46:02 2026 -0700
Deprecate deferrable implying wait in no-wait-default AWS operators (#75361)
These operators default to ``wait_for_completion=False``, but
``deferrable=True`` has always made them defer and wait anyway, and the
docstrings promised that deferring implies waiting. With
``[operators] default_deferrable`` enabled, every task using them waits
even though none asked to. The next major release will let
``wait_for_completion`` decide whether they wait, so ``deferrable=True``
on its own will stop waiting. Warning now, including when
``wait_for_completion`` was left at its default, gives Dag authors a
release to set it explicitly.
---
.../airflow/providers/amazon/aws/operators/eks.py | 68 +++++++------
.../airflow/providers/amazon/aws/operators/emr.py | 13 ++-
.../amazon/aws/operators/lambda_function.py | 7 +-
.../airflow/providers/amazon/aws/operators/mwaa.py | 8 +-
.../amazon/aws/operators/redshift_cluster.py | 9 +-
.../providers/amazon/aws/operators/sagemaker.py | 9 +-
.../tests/unit/amazon/aws/operators/test_eks.py | 111 +++++++++++++++------
.../amazon/aws/operators/test_emr_add_steps.py | 51 +++++-----
.../amazon/aws/operators/test_lambda_function.py | 10 +-
.../tests/unit/amazon/aws/operators/test_mwaa.py | 11 +-
.../amazon/aws/operators/test_redshift_cluster.py | 22 +++-
.../aws/operators/test_sagemaker_pipeline.py | 21 +++-
12 files changed, 233 insertions(+), 107 deletions(-)
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
index 4840c1a38a0..9f8be514251 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/eks.py
@@ -47,6 +47,7 @@ from airflow.providers.amazon.aws.utils import (
build_resource_in_use_retry_args,
validate_execute_complete_event,
)
+from airflow.providers.amazon.aws.utils.deferral import should_defer
from airflow.providers.amazon.aws.utils.mixins import AwsHookParams,
aws_template_fields
from airflow.providers.amazon.aws.utils.waiter_with_logging import wait
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
@@ -268,8 +269,9 @@ class EksCreateClusterOperator(AwsBaseOperator[EksHook]):
:param waiter_delay: Time (in seconds) to wait between two consecutive
calls to check cluster state
:param waiter_max_attempts: The maximum number of attempts to check
cluster state
:param deferrable: If True, the operator will wait asynchronously for the
job to complete.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
:param delete_nodegroup_on_failure: Whether to attempt best-effort
deletion of the managed nodegroup if creation
fails during the wait phase after successful initiation. Defaults to
True.
"""
@@ -321,8 +323,6 @@ class EksCreateClusterOperator(AwsBaseOperator[EksHook]):
self.nodegroup_role_arn = nodegroup_role_arn
self.fargate_pod_execution_role_arn = fargate_pod_execution_role_arn
self.create_fargate_profile_kwargs = create_fargate_profile_kwargs or
{}
- if deferrable:
- wait_for_completion = False
self.wait_for_completion = wait_for_completion
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
@@ -366,7 +366,11 @@ class EksCreateClusterOperator(AwsBaseOperator[EksHook]):
# Short circuit early if we don't need to wait to attach compute
# and the caller hasn't requested to wait for the cluster either.
- if not any([self.compute, self.wait_for_completion, self.deferrable]):
+ if (
+ not self.compute
+ and not self.wait_for_completion
+ and not should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion)
+ ):
return None
self.log.info("Waiting for EKS Cluster to provision. This will take
some time.")
@@ -446,7 +450,7 @@ class EksCreateClusterOperator(AwsBaseOperator[EksHook]):
cluster_name=self.cluster_name,
aws_conn_id=self.aws_conn_id,
region=self.region_name,
- wait_for_completion=self.wait_for_completion,
+ wait_for_completion=False,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
nodegroup_name=self.nodegroup_name,
@@ -459,6 +463,11 @@ class EksCreateClusterOperator(AwsBaseOperator[EksHook]):
subnets=cast("list[str]",
self.resources_vpc_config.get("subnetIds")),
delete_nodegroup_on_failure=self.delete_nodegroup_on_failure,
)
+ # The cluster wait above is needed to attach compute; waiting for
the compute itself is not.
+ if not self.compute or not should_defer(
+ deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion
+ ):
+ return
if self.compute == "fargate":
self.defer(
trigger=EksCreateFargateProfileTrigger(
@@ -535,8 +544,9 @@ class EksCreateNodegroupOperator(AwsBaseOperator[EksHook]):
:param waiter_delay: Time (in seconds) to wait between two consecutive
calls to check nodegroup state
:param waiter_max_attempts: The maximum number of attempts to check
nodegroup state
:param deferrable: If True, the operator will wait asynchronously for the
nodegroup to be created.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
:param delete_nodegroup_on_failure: Whether to attempt best-effort
deletion of the managed nodegroup if creation
fails during the wait phase after successful initiation. Defaults to
True.
"""
@@ -572,8 +582,6 @@ class EksCreateNodegroupOperator(AwsBaseOperator[EksHook]):
self.nodegroup_role_arn = nodegroup_role_arn
self.nodegroup_name = nodegroup_name
self.create_nodegroup_kwargs = create_nodegroup_kwargs or {}
- if deferrable:
- wait_for_completion = False
self.wait_for_completion = wait_for_completion
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
@@ -605,12 +613,13 @@ class
EksCreateNodegroupOperator(AwsBaseOperator[EksHook]):
)
self.nodegroup_subnets = nodegroup_subnets_list
+ defer = should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion)
_create_compute(
compute=self.compute,
cluster_name=self.cluster_name,
aws_conn_id=self.aws_conn_id,
region=self.region_name,
- wait_for_completion=self.wait_for_completion,
+ wait_for_completion=self.wait_for_completion and not defer,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
nodegroup_name=self.nodegroup_name,
@@ -620,7 +629,7 @@ class EksCreateNodegroupOperator(AwsBaseOperator[EksHook]):
delete_nodegroup_on_failure=self.delete_nodegroup_on_failure,
)
- if self.deferrable:
+ if defer:
self.defer(
trigger=EksCreateNodegroupTrigger(
cluster_name=self.cluster_name,
@@ -672,8 +681,9 @@ class
EksCreateFargateProfileOperator(AwsBaseOperator[EksHook]):
:param waiter_delay: Time (in seconds) to wait between two consecutive
calls to check profile status
:param waiter_max_attempts: The maximum number of attempts to check the
status of the profile.
:param deferrable: If True, the operator will wait asynchronously for the
profile to be created.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
"""
aws_hook_class = EksHook
@@ -705,8 +715,6 @@ class
EksCreateFargateProfileOperator(AwsBaseOperator[EksHook]):
self.pod_execution_role_arn = pod_execution_role_arn
self.fargate_profile_name = fargate_profile_name
self.create_fargate_profile_kwargs = create_fargate_profile_kwargs or
{}
- if deferrable:
- wait_for_completion = False
self.wait_for_completion = wait_for_completion
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
@@ -723,12 +731,13 @@ class
EksCreateFargateProfileOperator(AwsBaseOperator[EksHook]):
super().__init__(**kwargs)
def execute(self, context: Context):
+ defer = should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion)
_create_compute(
compute=self.compute,
cluster_name=self.cluster_name,
aws_conn_id=self.aws_conn_id,
region=self.region_name,
- wait_for_completion=self.wait_for_completion,
+ wait_for_completion=self.wait_for_completion and not defer,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
fargate_profile_name=self.fargate_profile_name,
@@ -736,7 +745,7 @@ class
EksCreateFargateProfileOperator(AwsBaseOperator[EksHook]):
fargate_selectors=self.selectors,
create_fargate_profile_kwargs=self.create_fargate_profile_kwargs,
)
- if self.deferrable:
+ if defer:
self.defer(
trigger=EksCreateFargateProfileTrigger(
cluster_name=self.cluster_name,
@@ -787,8 +796,9 @@ class EksDeleteClusterOperator(AwsBaseOperator[EksHook]):
:param waiter_delay: Time (in seconds) to wait between two consecutive
calls to check cluster state
:param waiter_max_attempts: The maximum number of attempts to check
cluster state
:param deferrable: If True, the operator will wait asynchronously for the
cluster to be deleted.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
"""
@@ -810,8 +820,6 @@ class EksDeleteClusterOperator(AwsBaseOperator[EksHook]):
) -> None:
self.cluster_name = cluster_name
self.force_delete_compute = force_delete_compute
- if deferrable:
- wait_for_completion = False
self.wait_for_completion = wait_for_completion
self.deferrable = deferrable
self.waiter_delay = waiter_delay
@@ -828,7 +836,7 @@ class EksDeleteClusterOperator(AwsBaseOperator[EksHook]):
super().__init__(**kwargs)
def execute(self, context: Context):
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=EksDeleteClusterTrigger(
cluster_name=self.cluster_name,
@@ -941,8 +949,9 @@ class EksDeleteNodegroupOperator(AwsBaseOperator[EksHook]):
:param waiter_delay: Time (in seconds) to wait between two consecutive
calls to check nodegroup state
:param waiter_max_attempts: The maximum number of attempts to check
nodegroup state
:param deferrable: If True, the operator will wait asynchronously for the
nodegroup to be deleted.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
"""
@@ -982,7 +991,7 @@ class EksDeleteNodegroupOperator(AwsBaseOperator[EksHook]):
for attempt in
tenacity.Retrying(**build_resource_in_use_retry_args(self.log)):
with attempt:
self.hook.delete_nodegroup(clusterName=self.cluster_name,
nodegroupName=self.nodegroup_name)
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=EksDeleteNodegroupTrigger(
cluster_name=self.cluster_name,
@@ -1037,8 +1046,9 @@ class
EksDeleteFargateProfileOperator(AwsBaseOperator[EksHook]):
:param waiter_delay: Time (in seconds) to wait between two consecutive
calls to check profile status
:param waiter_max_attempts: The maximum number of attempts to check the
status of the profile.
:param deferrable: If True, the operator will wait asynchronously for the
profile to be deleted.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
"""
aws_hook_class = EksHook
@@ -1078,7 +1088,7 @@ class
EksDeleteFargateProfileOperator(AwsBaseOperator[EksHook]):
self.hook.delete_fargate_profile(
clusterName=self.cluster_name,
fargateProfileName=self.fargate_profile_name
)
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=EksDeleteFargateProfileTrigger(
cluster_name=self.cluster_name,
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py
index 5c2356db030..ef4bd472076 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/emr.py
@@ -53,6 +53,7 @@ from airflow.providers.amazon.aws.triggers.emr import (
EmrTerminateJobFlowTrigger,
)
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.deferral import should_defer
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.amazon.aws.utils.waiter import (
WAITER_POLICY_NAME_MAPPING,
@@ -102,8 +103,9 @@ class EmrAddStepsOperator(AwsBaseOperator[EmrHook]):
:param execution_role_arn: The ARN of the runtime role for a step on the
cluster.
:param do_xcom_push: if True, job_flow_id is pushed to XCom with key
job_flow_id.
:param deferrable: If True, the operator will wait asynchronously for the
job to complete.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
+ This mode requires aiobotocore module to be installed. (default: False)
:param openlineage_inject_parent_job_info: If True, injects OpenLineage
parent job information
into Spark steps so the Spark job emits a ``parentRunFacet`` linking
back to the Airflow task.
Defaults to the ``openlineage.spark_inject_parent_job_info`` config
value.
@@ -151,7 +153,7 @@ class EmrAddStepsOperator(AwsBaseOperator[EmrHook]):
self.job_flow_name = job_flow_name
self.cluster_states = cluster_states
self.steps = steps
- self.wait_for_completion = False if deferrable else wait_for_completion
+ self.wait_for_completion = wait_for_completion
self.waiter_delay = waiter_delay
self.waiter_max_attempts = waiter_max_attempts
self.execution_role_arn = execution_role_arn
@@ -224,15 +226,16 @@ class EmrAddStepsOperator(AwsBaseOperator[EmrHook]):
if self.openlineage_inject_parent_job_info:
self.log.info("Injecting OpenLineage parent job information into
EMR Spark steps.")
steps = self._inject_openlineage_parent_job_information(steps,
context)
+ defer = should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion)
step_ids = self.hook.add_job_flow_steps(
job_flow_id=job_flow_id,
steps=steps,
- wait_for_completion=self.wait_for_completion,
+ wait_for_completion=self.wait_for_completion and not defer,
waiter_delay=self.waiter_delay,
waiter_max_attempts=self.waiter_max_attempts,
execution_role_arn=self.execution_role_arn,
)
- if self.deferrable:
+ if defer:
self.defer(
trigger=EmrAddStepsTrigger(
job_flow_id=job_flow_id,
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/lambda_function.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/lambda_function.py
index 28da2e244e5..201fa069554 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/operators/lambda_function.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/operators/lambda_function.py
@@ -26,6 +26,7 @@ from airflow.providers.amazon.aws.hooks.lambda_function
import LambdaHook
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.triggers.lambda_function import
LambdaCreateFunctionCompleteTrigger
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.deferral import should_defer
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.common.compat.sdk import AirflowException, conf
@@ -58,8 +59,10 @@ class
LambdaCreateFunctionOperator(AwsBaseOperator[LambdaHook]):
:param waiter_max_attempts: Maximum number of attempts to poll the
creation.
:param waiter_delay: Number of seconds between polling the state of the
creation.
:param deferrable: If True, the operator will wait asynchronously for the
creation to complete.
- This implies waiting for creation complete. This mode requires
aiobotocore module to be installed.
+ This mode requires aiobotocore module to be installed.
(default: False, but can be overridden in config file by setting
default_deferrable to True)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
:param aws_conn_id: The AWS connection ID to use
"""
@@ -119,7 +122,7 @@ class
LambdaCreateFunctionOperator(AwsBaseOperator[LambdaHook]):
)
self.log.info("Lambda response: %r", response)
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=LambdaCreateFunctionCompleteTrigger(
function_name=self.function_name,
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py
index ee468f7a4cd..3225343ff68 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/mwaa.py
@@ -25,6 +25,7 @@ from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.triggers.mwaa import
MwaaDagRunCompletedTrigger
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.deferral import should_defer
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.common.compat.sdk import AirflowException, conf
@@ -58,8 +59,9 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
:param waiter_delay: Time in seconds to wait between status checks.
(default: 60)
:param waiter_max_attempts: Maximum number of attempts to check for Dag
run completion. (default: 20)
:param deferrable: If True, the operator will wait asynchronously for the
Dag run to stop.
- This implies waiting for completion. This mode requires aiobotocore
module to be installed.
- (default: False)
+ This mode requires aiobotocore module to be installed. (default: False)
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is 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
@@ -160,7 +162,7 @@ class MwaaTriggerDagRunOperator(AwsBaseOperator[MwaaHook]):
self.log.info("Dag run %s of Dag %s created", dag_run_id,
self.trigger_dag_id)
task_description = f"Dag run {dag_run_id} of Dag {self.trigger_dag_id}
to complete"
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.log.info("Deferring for %s", task_description)
self.defer(
trigger=MwaaDagRunCompletedTrigger(
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py
index 2424e35e0ca..942bea66f13 100644
---
a/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py
+++
b/providers/amazon/src/airflow/providers/amazon/aws/operators/redshift_cluster.py
@@ -35,6 +35,7 @@ from airflow.providers.amazon.aws.triggers.redshift_cluster
import (
RedshiftResumeClusterTrigger,
)
from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.deferral import should_defer
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.common.compat.sdk import AirflowException, conf
from airflow.utils.helpers import prune_dict
@@ -110,6 +111,8 @@ class
RedshiftCreateClusterOperator(AwsBaseOperator[RedshiftHook]):
:param max_attempt: The maximum number of attempts to be made. Default: 5
:param poll_interval: The amount of time in seconds to wait between
attempts. Default: 60
:param deferrable: If True, the operator will run in deferrable mode.
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
:param delete_cluster_on_failure: If True, best-effort deletion of the
redshift cluster will be attempted
after post-creation failure. Default: True.
:param cleanup_timeout_seconds: Maximum time in seconds to attempt
@@ -372,7 +375,7 @@ class
RedshiftCreateClusterOperator(AwsBaseOperator[RedshiftHook]):
self.master_user_password,
params,
)
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=RedshiftCreateClusterTrigger(
cluster_identifier=self.cluster_identifier,
@@ -448,6 +451,8 @@ class
RedshiftCreateClusterSnapshotOperator(AwsBaseOperator[RedshiftHook]):
:param verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param deferrable: If True, the operator will run as a deferrable operator.
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
"""
template_fields: Sequence[str] = aws_template_fields(
@@ -494,7 +499,7 @@ class
RedshiftCreateClusterSnapshotOperator(AwsBaseOperator[RedshiftHook]):
retention_period=self.retention_period,
tags=self.tags,
)
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=RedshiftCreateClusterSnapshotTrigger(
cluster_identifier=self.cluster_identifier,
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
index d8c7d0aa06b..a7a58e9bf5a 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/sagemaker.py
@@ -38,6 +38,7 @@ from airflow.providers.amazon.aws.triggers.sagemaker import (
SageMakerTrigger,
)
from airflow.providers.amazon.aws.utils import trim_none_values,
validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.deferral import should_defer
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.providers.amazon.aws.utils.sagemaker import ApprovalStatus
from airflow.providers.amazon.aws.utils.tags import format_tags
@@ -1443,6 +1444,8 @@ class
SageMakerStartPipelineOperator(SageMakerBaseOperator):
:param verbose: Whether to print steps details when waiting for completion.
Defaults to true, consider turning off for pipelines that have
thousands of steps.
:param deferrable: Run operator in the deferrable mode.
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
:return str: Returns The ARN of the pipeline execution created in Amazon
SageMaker.
"""
@@ -1485,7 +1488,7 @@ class
SageMakerStartPipelineOperator(SageMakerBaseOperator):
self.log.info(
"Starting a new execution for pipeline %s, running with ARN %s",
self.pipeline_name, arn
)
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=SageMakerPipelineTrigger(
waiter_type=SageMakerPipelineTrigger.Type.COMPLETE,
@@ -1542,6 +1545,8 @@ class
SageMakerStopPipelineOperator(SageMakerBaseOperator):
Defaults to true, consider turning off for pipelines that have
thousands of steps.
:param fail_if_not_running: raises an exception if the pipeline stopped or
succeeded before this was run
:param deferrable: Run operator in the deferrable mode.
+ Deferring when ``wait_for_completion`` is False is deprecated; from
the next major release
+ this has no effect when ``wait_for_completion`` is False.
:return str: Returns the status of the pipeline execution after the
operation has been done.
"""
@@ -1587,7 +1592,7 @@ class
SageMakerStopPipelineOperator(SageMakerBaseOperator):
return status
# else, eventually wait for completion
- if self.deferrable:
+ if should_defer(deferrable=self.deferrable,
wait_for_completion=self.wait_for_completion):
self.defer(
trigger=SageMakerPipelineTrigger(
waiter_type=SageMakerPipelineTrigger.Type.STOPPED,
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
index 418878c777c..e76eb351f4c 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_eks.py
@@ -43,6 +43,7 @@ from airflow.providers.amazon.aws.triggers.eks import (
EksCreateFargateProfileTrigger,
EksCreateNodegroupTrigger,
EksDeleteFargateProfileTrigger,
+ EksDeleteNodegroupTrigger,
EksPodTrigger,
)
from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction
@@ -58,6 +59,7 @@ from unit.amazon.aws.utils.eks_test_constants import (
TASK_ID,
)
from unit.amazon.aws.utils.eks_test_utils import convert_keys
+from unit.amazon.aws.utils.test_deferral import expect_deferral_warning
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
from unit.amazon.aws.utils.test_waiter import assert_expected_waiter_type
@@ -377,21 +379,53 @@ class TestEksCreateClusterOperator:
eks_create_cluster_operator.execute({})
assert len(caplog.records) == 0
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
@mock.patch.object(EksHook, "create_cluster")
- def test_eks_create_cluster_with_deferrable(self, mock_create_cluster,
caplog):
+ def test_eks_create_cluster_with_deferrable(self, mock_create_cluster,
caplog, wait_for_completion):
mock_create_cluster.return_value = None
eks_create_cluster_operator = EksCreateClusterOperator(
task_id=TASK_ID,
**self.create_cluster_params,
compute=None,
- wait_for_completion=False,
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred):
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred),
+ ):
eks_create_cluster_operator.execute({})
assert "Waiting for EKS Cluster to provision. This will take some
time." in caplog.messages
+ @pytest.mark.parametrize(
+ ("compute", "trigger_class"),
+ [
+ pytest.param("nodegroup", EksCreateNodegroupTrigger,
id="nodegroup"),
+ pytest.param("fargate", EksCreateFargateProfileTrigger,
id="fargate"),
+ ],
+ )
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
+ @mock.patch("airflow.providers.amazon.aws.operators.eks._create_compute",
autospec=True)
+ def test_deferrable_create_cluster_next_defers_for_compute(
+ self, mock_create_compute, wait_for_completion, compute, trigger_class
+ ):
+ operator = EksCreateClusterOperator(
+ task_id=TASK_ID,
+ **self.create_cluster_params,
+ compute=compute,
+ wait_for_completion=wait_for_completion,
+ deferrable=True,
+ )
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as exc,
+ ):
+ operator.deferrable_create_cluster_next({}, event={"status":
"success"})
+
+ assert isinstance(exc.value.trigger, trigger_class)
+ assert mock_create_compute.call_args.kwargs["wait_for_completion"] is
False
+
def test_template_fields(self):
op = EksCreateClusterOperator(
task_id=TASK_ID, **self.create_cluster_params, compute="fargate",
region_name="us-east-1"
@@ -465,16 +499,23 @@ class TestEksCreateFargateProfileOperator:
)
assert_expected_waiter_type(mock_waiter, "FargateProfileActive")
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
+ @mock.patch.object(Waiter, "wait")
@mock.patch.object(EksHook, "create_fargate_profile")
- def test_create_fargate_profile_deferrable(self, _):
+ def test_create_fargate_profile_deferrable(self, _, mock_waiter,
wait_for_completion):
op_kwargs = {**self.create_fargate_profile_params}
operator = EksCreateFargateProfileOperator(
task_id=TASK_ID,
**op_kwargs,
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred) as exc:
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as exc,
+ ):
operator.execute({})
+ mock_waiter.assert_not_called()
assert isinstance(exc.value.trigger, EksCreateFargateProfileTrigger), (
"Trigger is not a EksCreateFargateProfileTrigger"
)
@@ -559,35 +600,26 @@ class TestEksCreateNodegroupOperator:
)
assert_expected_waiter_type(mock_waiter, "NodegroupActive")
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
+ @mock.patch.object(Waiter, "wait")
@mock.patch.object(EksHook, "create_nodegroup")
- def test_create_nodegroup_deferrable(self, mock_create_nodegroup):
+ def test_create_nodegroup_deferrable(self, mock_create_nodegroup,
mock_waiter, wait_for_completion):
mock_create_nodegroup.return_value = True
op_kwargs = {**self.create_nodegroup_params}
operator = EksCreateNodegroupOperator(
task_id=TASK_ID,
**op_kwargs,
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred) as exc:
+ assert operator.wait_for_completion is wait_for_completion
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as exc,
+ ):
operator.execute({})
assert isinstance(exc.value.trigger, EksCreateNodegroupTrigger)
-
- def test_create_nodegroup_deferrable_versus_wait_for_completion(self):
- op_kwargs = {**self.create_nodegroup_params}
- operator = EksCreateNodegroupOperator(
- task_id=TASK_ID,
- **op_kwargs,
- deferrable=True,
- wait_for_completion=True,
- )
- assert operator.wait_for_completion is False
- operator = EksCreateNodegroupOperator(
- task_id=TASK_ID,
- **op_kwargs,
- deferrable=False,
- wait_for_completion=True,
- )
- assert operator.wait_for_completion is True
+ mock_waiter.assert_not_called()
def test_template_fields(self):
op_kwargs = {**self.create_nodegroup_params}
@@ -763,9 +795,14 @@ class TestEksDeleteClusterOperator:
)
assert_expected_waiter_type(mock_waiter, "ClusterDeleted")
- def test_eks_delete_cluster_operator_with_deferrable(self):
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
+ def test_eks_delete_cluster_operator_with_deferrable(self,
wait_for_completion):
self.delete_cluster_operator.deferrable = True
- with pytest.raises(TaskDeferred):
+ self.delete_cluster_operator.wait_for_completion = wait_for_completion
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred),
+ ):
self.delete_cluster_operator.execute({})
@pytest.mark.parametrize("status", ["deleted", "success"])
@@ -975,6 +1012,19 @@ class TestEksDeleteNodegroupOperator:
)
assert_expected_waiter_type(mock_waiter, "NodegroupDeleted")
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
+ @mock.patch.object(EksHook, "delete_nodegroup")
+ def test_delete_nodegroup_deferrable(self, _, wait_for_completion):
+ self.delete_nodegroup_operator.deferrable = True
+ self.delete_nodegroup_operator.wait_for_completion =
wait_for_completion
+
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as exc,
+ ):
+ self.delete_nodegroup_operator.execute({})
+ assert isinstance(exc.value.trigger, EksDeleteNodegroupTrigger)
+
def test_template_fields(self):
validate_template_fields(self.delete_nodegroup_operator)
@@ -1026,11 +1076,16 @@ class TestEksDeleteFargateProfileOperator:
)
assert_expected_waiter_type(mock_waiter, "FargateProfileDeleted")
+ @pytest.mark.parametrize("wait_for_completion", [True, False],
ids=["wait", "no_wait"])
@mock.patch.object(EksHook, "delete_fargate_profile")
- def test_delete_fargate_profile_deferrable(self, _):
+ def test_delete_fargate_profile_deferrable(self, _, wait_for_completion):
self.delete_fargate_profile_operator.deferrable = True
+ self.delete_fargate_profile_operator.wait_for_completion =
wait_for_completion
- with pytest.raises(TaskDeferred) as exc:
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as exc,
+ ):
self.delete_fargate_profile_operator.execute({})
assert isinstance(exc.value.trigger, EksDeleteFargateProfileTrigger), (
"Trigger is not a EksDeleteFargateProfileTrigger"
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
index 7a8d6b3892e..2a45a1a0931 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_emr_add_steps.py
@@ -36,6 +36,7 @@ from tests_common.test_utils.compat import timezone
from tests_common.test_utils.dag import sync_dag_to_db
from tests_common.test_utils.taskinstance import create_task_instance,
render_template_fields
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+from unit.amazon.aws.utils.test_deferral import expect_deferral_warning
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
DEFAULT_DATE = timezone.datetime(2017, 1, 1)
@@ -282,38 +283,42 @@ class TestEmrAddStepsOperator:
execution_role_arn=None,
)
- def test_wait_for_completion_false_with_deferrable(self):
- job_flow_id = "j-8989898989"
- operator = EmrAddStepsOperator(
- task_id="test_task",
- job_flow_id=job_flow_id,
- aws_conn_id="aws_default",
- dag=DAG("test_dag_id", schedule=None, default_args=self.args),
- wait_for_completion=True,
- deferrable=True,
- )
-
- assert operator.wait_for_completion is False
-
+ @pytest.mark.parametrize(
+ ("wait_for_completion", "deferrable"),
+ [
+ pytest.param(False, False, id="no_wait"),
+ pytest.param(True, False, id="wait"),
+ pytest.param(True, True, id="defer"),
+ pytest.param(False, True, id="defer_overrides_no_wait"),
+ ],
+ )
@patch("airflow.providers.amazon.aws.operators.emr.get_log_uri")
@patch("airflow.providers.amazon.aws.hooks.emr.EmrHook.add_job_flow_steps")
- def test_emr_add_steps_deferrable(self, mock_add_job_flow_steps,
mock_get_log_uri):
- mock_add_job_flow_steps.return_value = "test_step_id"
+ def test_emr_add_steps_wait_combinations(
+ self, mock_add_job_flow_steps, mock_get_log_uri, wait_for_completion,
deferrable
+ ):
+ mock_add_job_flow_steps.return_value = ["test_step_id"]
mock_get_log_uri.return_value = "test/log/uri"
- job_flow_id = "j-8989898989"
operator = EmrAddStepsOperator(
task_id="test_task",
- job_flow_id=job_flow_id,
+ job_flow_id="j-8989898989",
aws_conn_id="aws_default",
dag=DAG("test_dag_id", schedule=None, default_args=self.args),
- wait_for_completion=True,
- deferrable=True,
+ wait_for_completion=wait_for_completion,
+ deferrable=deferrable,
)
+ assert operator.wait_for_completion is wait_for_completion
+
+ with expect_deferral_warning(deferrable=deferrable,
wait_for_completion=wait_for_completion):
+ if deferrable:
+ with pytest.raises(TaskDeferred) as exc:
+ operator.execute(MagicMock())
+ assert isinstance(exc.value.trigger, EmrAddStepsTrigger)
+ else:
+ operator.execute(MagicMock())
- with pytest.raises(TaskDeferred) as exc:
- operator.execute(MagicMock())
-
- assert isinstance(exc.value.trigger, EmrAddStepsTrigger), "Trigger is
not a EmrAddStepsTrigger"
+ hook_waits =
mock_add_job_flow_steps.call_args.kwargs["wait_for_completion"]
+ assert hook_waits is (wait_for_completion and not deferrable)
def test_template_fields(self):
op = EmrAddStepsOperator(
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_lambda_function.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_lambda_function.py
index e69e8867234..e378dc28842 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_lambda_function.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_lambda_function.py
@@ -30,6 +30,7 @@ from airflow.providers.amazon.aws.operators.lambda_function
import (
)
from airflow.providers.common.compat.sdk import TaskDeferred
+from unit.amazon.aws.utils.test_deferral import expect_deferral_warning
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
FUNCTION_NAME = "function_name"
@@ -114,8 +115,9 @@ class TestLambdaCreateFunctionOperator:
mock_hook_create_lambda.assert_called_once()
mock_hook_conn.get_waiter.assert_called_once_with("function_active_v2")
+ @pytest.mark.parametrize("wait_for_completion", [True, False])
@mock.patch.object(LambdaHook, "create_lambda")
- def test_create_lambda_deferrable(self, _):
+ def test_create_lambda_deferrable(self, _, wait_for_completion):
operator = LambdaCreateFunctionOperator(
task_id="task_test",
function_name=FUNCTION_NAME,
@@ -123,9 +125,13 @@ class TestLambdaCreateFunctionOperator:
code={
"ImageUri": IMAGE_URI,
},
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred):
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred),
+ ):
operator.execute(None)
@mock.patch.object(LambdaHook, "create_lambda")
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py
index 4e91ab1e1b4..59dbe1af139 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_mwaa.py
@@ -25,6 +25,7 @@ from airflow.providers.amazon.aws.hooks.mwaa import MwaaHook
from airflow.providers.amazon.aws.operators.mwaa import
MwaaTriggerDagRunOperator
from airflow.providers.common.compat.sdk import TaskDeferred
+from unit.amazon.aws.utils.test_deferral import expect_deferral_warning
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
OP_KWARGS = {
@@ -102,7 +103,7 @@ class TestMwaaTriggerDagRunOperator:
[
pytest.param(False, False, id="no_wait"),
pytest.param(True, False, id="wait"),
- pytest.param(False, True, id="defer"),
+ pytest.param(False, True, id="defer_overrides_no_wait"),
],
)
@mock.patch.object(MwaaHook, "get_waiter")
@@ -114,7 +115,8 @@ class TestMwaaTriggerDagRunOperator:
op = MwaaTriggerDagRunOperator(**OP_KWARGS)
mock_hook.invoke_rest_api.return_value = HOOK_RETURN_VALUE
op.defer = MagicMock()
- response = op.execute({})
+ with expect_deferral_warning(deferrable=deferrable,
wait_for_completion=wait_for_completion):
+ response = op.execute({})
assert response == HOOK_RETURN_VALUE
assert mock_hook.get_waiter.call_count == wait_for_completion
@@ -130,7 +132,10 @@ class TestMwaaTriggerDagRunOperator:
botocore_config=BOTOCORE_CONFIG,
)
- with pytest.raises(TaskDeferred) as exc_info:
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=False),
+ pytest.raises(TaskDeferred) as exc_info,
+ ):
op.execute({})
assert exc_info.value.trigger.serialize()[1] == {
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py
index 46d25b8ab33..acd1fda5faf 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_redshift_cluster.py
@@ -42,6 +42,7 @@ from airflow.providers.amazon.aws.triggers.redshift_cluster
import (
)
from airflow.providers.common.compat.sdk import AirflowException, TaskDeferred
+from unit.amazon.aws.utils.test_deferral import expect_deferral_warning
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
@@ -127,8 +128,9 @@ class TestRedshiftCreateClusterOperator:
# wait_for_completion is False so check waiter is not called
mock_conn.get_waiter.assert_not_called()
+ @pytest.mark.parametrize("wait_for_completion", [True, False])
@mock.patch.object(RedshiftHook, "conn")
- def test_create_cluster_deferrable(self, mock_conn):
+ def test_create_cluster_deferrable(self, mock_conn, wait_for_completion):
redshift_operator = RedshiftCreateClusterOperator(
task_id="task_test",
cluster_identifier="test-cluster",
@@ -136,11 +138,16 @@ class TestRedshiftCreateClusterOperator:
master_username="adminuser",
master_user_password="Test123$",
cluster_type="single-node",
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred):
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred),
+ ):
redshift_operator.execute(None)
+ mock_conn.get_waiter.assert_not_called()
def test_template_fields(self):
operator = RedshiftCreateClusterOperator(
@@ -369,18 +376,25 @@ class TestRedshiftCreateClusterSnapshotOperator:
WaiterConfig={"Delay": 15, "MaxAttempts": 20},
)
+ @pytest.mark.parametrize("wait_for_completion", [True, False])
@mock.patch.object(RedshiftHook, "cluster_status")
@mock.patch.object(RedshiftHook, "create_cluster_snapshot")
- def test_create_cluster_snapshot_deferred(self,
mock_create_cluster_snapshot, mock_cluster_status):
+ def test_create_cluster_snapshot_deferred(
+ self, mock_create_cluster_snapshot, mock_cluster_status,
wait_for_completion
+ ):
mock_cluster_status.return_value = "available"
mock_create_cluster_snapshot.return_value = True
create_snapshot = RedshiftCreateClusterSnapshotOperator(
task_id="test_snapshot",
cluster_identifier="test_cluster",
snapshot_identifier="test_snapshot",
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred) as exc:
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as exc,
+ ):
create_snapshot.execute(None)
assert isinstance(exc.value.trigger,
RedshiftCreateClusterSnapshotTrigger), (
"Trigger is not a RedshiftCreateClusterSnapshotTrigger"
diff --git
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_pipeline.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_pipeline.py
index fbfd2baafba..2423db51d66 100644
---
a/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_pipeline.py
+++
b/providers/amazon/tests/unit/amazon/aws/operators/test_sagemaker_pipeline.py
@@ -30,6 +30,7 @@ from airflow.providers.amazon.aws.operators.sagemaker import (
from airflow.providers.amazon.aws.triggers.sagemaker import
SageMakerPipelineTrigger
from airflow.providers.common.compat.sdk import TaskDeferred
+from unit.amazon.aws.utils.test_deferral import expect_deferral_warning
from unit.amazon.aws.utils.test_template_fields import validate_template_fields
if TYPE_CHECKING:
@@ -64,15 +65,20 @@ class TestSageMakerStartPipelineOperator:
)
check_status.assert_called_once()
+ @pytest.mark.parametrize("wait_for_completion", [True, False])
@mock.patch.object(SageMakerHook, "start_pipeline")
- def test_defer(self, start_mock):
+ def test_defer(self, start_mock, wait_for_completion):
op = SageMakerStartPipelineOperator(
task_id="test_sagemaker_operator",
pipeline_name="my_pipeline",
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred) as defer:
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as defer,
+ ):
op.execute({})
assert isinstance(defer.value.trigger, SageMakerPipelineTrigger)
@@ -83,6 +89,7 @@ class TestSageMakerStartPipelineOperator:
op = SageMakerStartPipelineOperator(
task_id="test_sagemaker_operator",
pipeline_name="my_pipeline",
+ wait_for_completion=True,
deferrable=True,
region_name=REGION_NAME,
verify=VERIFY,
@@ -118,16 +125,21 @@ class TestSageMakerStopPipelineOperator:
fail_if_not_running=False,
)
+ @pytest.mark.parametrize("wait_for_completion", [True, False])
@mock.patch.object(SageMakerHook, "stop_pipeline")
- def test_defer(self, stop_mock: MagicMock):
+ def test_defer(self, stop_mock: MagicMock, wait_for_completion):
stop_mock.return_value = "Stopping"
op = SageMakerStopPipelineOperator(
task_id="test_sagemaker_operator",
pipeline_exec_arn="my_pipeline_arn",
+ wait_for_completion=wait_for_completion,
deferrable=True,
)
- with pytest.raises(TaskDeferred) as defer:
+ with (
+ expect_deferral_warning(deferrable=True,
wait_for_completion=wait_for_completion),
+ pytest.raises(TaskDeferred) as defer,
+ ):
op.execute({})
assert isinstance(defer.value.trigger, SageMakerPipelineTrigger)
@@ -139,6 +151,7 @@ class TestSageMakerStopPipelineOperator:
op = SageMakerStopPipelineOperator(
task_id="test_sagemaker_operator",
pipeline_exec_arn="my_pipeline_arn",
+ wait_for_completion=True,
deferrable=True,
region_name=REGION_NAME,
verify=VERIFY,
