如何通过Airflow轮询SageMaker Autopilot任务状态并解决重复通知问题

• Airflow中PythonSensor的每一次poke调用都是独立的运行实例，会重新加载DAG定义、执行状态检查函数，存在进程内存中的状态集合会在每次调用时被重新初始化，无法跨多次poke调用持久化。
• 本地简化测试用例是在同一个进程内连续执行模拟逻辑，内存状态可以持续保留，和Airflow的多进程/分布式运行模型存在本质差异，所以测试用例可以正常运行但正式环境失效。

不需要额外引入临时文件、独立数据库等外部存储，直接用Airflow自带的XCom组件存储已推送的历史状态即可，实现成本最低：

核心逻辑调整 #

修改_sagemaker_job_status函数，每次执行时先从XCom拉取已推送的状态集合，判断当前状态是否为新状态，完成推送后将更新后的状态集合写回XCom：

import boto3
import json

def _sagemaker_job_status(**context):
    ti = context["ti"]
    # 从XCom拉取历史已推送状态，无历史数据则初始化为空集合
    raw_pushed_status = ti.xcom_pull(
        key="sagemaker_pushed_status", 
        task_ids=ti.task_id
    )
    pushed_status_set = set(json.loads(raw_pushed_status)) if raw_pushed_status else set()
    
    # 获取上游传入的任务名，查询SageMaker任务状态
    automl_job_name = ti.xcom_pull(task_ids="你的上游SageMaker任务ID")
    sagemaker_client = boto3.client("sagemaker")
    job_detail = sagemaker_client.describe_auto_ml_job(AutoMLJobName=automl_job_name)
    current_status = (
        job_detail["AutoMLJobStatus"], 
        job_detail["AutoMLJobSecondaryStatus"]
    )

    # 仅新状态推送Slack
    if current_status not in pushed_status_set:
        # 此处替换为你的Slack推送逻辑
        send_slack_msg(f"SageMaker任务状态更新：{current_status[0]} - {current_status[1]}")
        pushed_status_set.add(current_status)
        # 序列化为JSON字符串存入XCom，兼容所有Airflow配置
        ti.xcom_push(key="sagemaker_pushed_status", value=json.dumps(list(pushed_status_set)))
    
    # 返回任务是否完成，控制Sensor是否终止
    return current_status[0] == "Completed" and current_status[1] == "Completed"

PythonSensor配置说明 #

确保你的PythonSensor开启了provide_context=True（Airflow 2.x默认开启）即可正常运行，不需要修改其他配置。

你提到的两个方案也可以实现，但是有额外限制：

• 临时JSON文件方案：如果使用多Worker部署，必须将文件存在S3等共享存储中，否则下次poke调度到其他Worker节点会读取不到历史状态
• 独立数据库存储方案：功能完全可行，但是需要额外维护状态表，仅适合状态流转非常复杂的场景使用

内容的提问来源于stack exchange，提问作者yudhiesh