如何在SageMaker Notebook中编程判断AWS Step Function执行完成？

核心结论 #

• boto3的Step Function客户端没有提供原生的阻塞式启动执行接口，所有官方客户端的start_execution调用本身都是异步返回的
• 可实现完全在Notebook Python环境内运行、不新增任何AWS资源、不修改原有Step Function定义的精准等待逻辑，比手动查列表、固定休眠的方案更可靠

方案核心优化点 #

• 放弃原有用list_executions(maxResults=1)查询状态的方式：该方法会因状态机上存在其他并发触发的执行任务，查到非本次启动的执行状态，出现判断偏差
• 直接绑定启动执行时返回的唯一executionArn作为查询锚点，调用describe_execution接口定向查询单条执行的实时状态，无串查风险
• 用轻量循环替代手动查询，一旦识别到执行进入终态立刻退出循环，衔接后续逻辑

可直接复用的代码实现 #

首先编写可复用的等待函数：

import time
def block_until_sfn_finish(sfn_client, target_execution_arn, check_interval=15):
    while True:
        resp = sfn_client.describe_execution(executionArn=target_execution_arn)
        status = resp["status"]
        print(f"轮询到执行状态: {status}")
        # 所有非运行中的状态都属于终态，直接退出循环
        if status in ["SUCCEEDED", "FAILED", "TIMED_OUT", "ABORTED"]:
            return resp
        time.sleep(check_interval)

完整流程衔接代码（原有启动执行逻辑可完全保留）：

state_machine_arn = 'arn:aws:states:us-west-1:1234567891:stateMachine:alexanderMyPackageStateMachineE3411O13-A1vQWERTP9q9'
sfn = boto3.client('stepfunctions')
# 原有kwargs传参逻辑保持不变
response = sfn.start_execution(**kwargs)
run_arn = response['executionArn']
print(f"Started run {run_name}. ARN is {run_arn}.")

# 自动等待执行完成
final_result = block_until_sfn_finish(sfn, run_arn, check_interval=30)

# 状态判断后衔接自定义后分析
if final_result["status"] == "SUCCEEDED":
    # 此处替换为自定义后分析逻辑
    print("执行成功，启动后分析流程")
else:
    print(f"执行异常终止，状态为{final_result['status']}，跳过后分析")

方案适配说明 #

• 全程不需要创建EventBridge规则、Lambda函数等额外AWS资源，所有逻辑在Notebook的Python进程内运行
• 不需要修改Step Function的状态定义，也不需要改动原有MyPackage的代码，完全匹配约束条件
• 轮询间隔可根据实际执行时长灵活调整：如果任务通常运行数小时，可将check_interval设为60秒，减少无意义的API调用；如果任务运行时长很短，可将间隔设为5秒，降低执行完成后的等待延迟
• 对比固定长休眠的方案，该逻辑从启动执行后就持续追踪目标任务状态，一旦执行完成会在一个轮询周期内触发后续逻辑，没有额外的无效等待

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