# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""
System test for SageMakerConditionOperator.
This operator evaluates conditions against XCom values passed from upstream tasks.
The DAG simulates an ML accuracy-gate workflow:
1. ``produce_metrics`` pushes a dict of metrics to XCom.
2. ``check_accuracy`` uses SageMakerConditionOperator to branch:
- accuracy >= 0.9 AND loss < 0.1 -> ``deploy_model``
- otherwise -> ``retrain_model``
3. Only the correct branch task runs; the other is skipped.
"""
from __future__ import annotations
from datetime import datetime
from airflow.providers.amazon.aws.operators.sagemaker import SageMakerConditionOperator
from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
if AIRFLOW_V_3_0_PLUS:
from airflow.sdk import DAG, chain, task
else:
from airflow.decorators import task # type: ignore[attr-defined,no-redef]
from airflow.models.baseoperator import chain # type: ignore[attr-defined,no-redef]
from airflow.models.dag import DAG # type: ignore[attr-defined,no-redef,assignment]
from system.amazon.aws.utils import SystemTestContextBuilder
DAG_ID = "example_sagemaker_condition"
sys_test_context_task = SystemTestContextBuilder().build()
with DAG(
DAG_ID,
schedule="@once",
start_date=datetime(2021, 1, 1),
catchup=False,
) as dag:
test_context = sys_test_context_task()
# TEST SETUP: push simulated ML metrics to XCom
@task
def produce_metrics():
"""Simulate an ML training job that returns accuracy and loss metrics."""
return {"accuracy": 0.95, "loss": 0.04}
metrics = produce_metrics()
# TEST BODY
# [START howto_operator_sagemaker_condition]
check_accuracy = SageMakerConditionOperator(
task_id="check_accuracy",
conditions=[
{
"type": "GreaterThanOrEqualTo",
"left_value": "{{ ti.xcom_pull(task_ids='produce_metrics')['accuracy'] }}",
"right_value": 0.9,
},
{
"type": "LessThan",
"left_value": "{{ ti.xcom_pull(task_ids='produce_metrics')['loss'] }}",
"right_value": 0.1,
},
],
if_task_ids=["deploy_model"],
else_task_ids=["retrain_model"],
)
# [END howto_operator_sagemaker_condition]
@task
def deploy_model():
"""Placeholder: model meets quality bar, proceed to deployment."""
return "deployed"
@task
def retrain_model():
"""Placeholder: model does not meet quality bar, retrain."""
return "retrained"
# Scenario 2: condition evaluates to False -> else branch
@task
def produce_bad_metrics():
"""Simulate a training job with poor accuracy."""
return {"accuracy": 0.5, "loss": 0.8}
bad_metrics = produce_bad_metrics()
# [START howto_operator_sagemaker_condition_flat]
check_bad_accuracy = SageMakerConditionOperator(
task_id="check_bad_accuracy",
condition_type="GreaterThanOrEqualTo",
left_value="{{ ti.xcom_pull(task_ids='produce_bad_metrics')['accuracy'] }}",
right_value=0.9,
if_task_ids=["should_not_run"],
else_task_ids=["should_run"],
)
# [END howto_operator_sagemaker_condition_flat]
@task
def should_not_run():
"""This task should be skipped because accuracy < 0.9."""
return "error: should not have run"
@task
def should_run():
"""This task should execute because accuracy < 0.9 -> else branch."""
return "correctly routed to else branch"
# Scenario 3: Or condition + Not condition
# [START howto_operator_sagemaker_condition_not_or]
check_logical = SageMakerConditionOperator(
task_id="check_logical",
conditions=[
{
"type": "Or",
"conditions": [
{"type": "Equals", "left_value": 1, "right_value": 2},
{"type": "Equals", "left_value": 3, "right_value": 3},
],
},
{
"type": "Not",
"condition": {"type": "Equals", "left_value": "a", "right_value": "b"},
},
],
if_task_ids=["logical_pass"],
else_task_ids=["logical_fail"],
)
# [END howto_operator_sagemaker_condition_not_or]
@task
def logical_pass():
"""Or(1==2, 3==3) -> True AND Not(a==b) -> True -> if branch."""
return "logical conditions passed"
@task
def logical_fail():
return "error: logical conditions should have passed"
test_context >> [metrics, bad_metrics, check_logical]
chain(metrics, check_accuracy, [deploy_model(), retrain_model()])
chain(bad_metrics, check_bad_accuracy, [should_not_run(), should_run()])
chain(check_logical, [logical_pass(), logical_fail()])
from tests_common.test_utils.watcher import watcher
# This test needs watcher in order to properly mark success/failure
# when "tearDown" task with trigger rule is part of the DAG
list(dag.tasks) >> watcher()
from tests_common.test_utils.system_tests import get_test_run # noqa: E402
# Needed to run the example DAG with pytest (see: contributing-docs/testing/system_tests.rst)
test_run = get_test_run(dag)