prefect resume pause

import json
import os
from datetime import datetime, timezone
from pathlib import Path
import time
from prefect import flow, task, apause_flow_run
from pydantic import BaseModel

class NextStepInput(BaseModel):
    next: str
    node: str = "default"


@task(persist_result=True) #
async def step1(message: str) -> dict[str, any]:
    time.sleep(2)  # Simulate some processing time
    print("step1 executed")
    return {
        "message": message,
        "created_at": datetime.now(timezone.utc).isoformat(),
        "steps":["step1"]
    }


@task(persist_result=True) #
async def step2(payload: dict[str, any]) -> str:
    time.sleep(2) 
    payload["steps"].append("step2")
    print("step2 executed")
    return payload

@task(persist_result=True) #
async def step3(payload: dict[str, any]) -> str:
    time.sleep(2) 
    payload["steps"].append("step3")
    print("step3 executed")
    # raise Exception("Simulated error in step3")  # 模拟异常，测试失败重试机制
    return payload


@task(persist_result=True) #
async def step4(payload: dict[str, any]) -> str:
    time.sleep(2) 
    payload["steps"].append("step4")
    print("step4 executed")
    output_dir = Path(os.getenv("DEMO_OUTPUT_DIR", "./runs"))
    output_dir.mkdir(parents=True, exist_ok=True)

    output_path = output_dir / f"prefect-demo-{payload['created_at'].replace(':', '-')}.json"
    output_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    return str(output_path)

@flow(name="map_pipeline_flow", log_prints=True)
async def map_pipeline_flow(message: str = "Hello from local Prefect") -> str:
    payload = await step1(message)

    # suspend_flow_run()
    res = await apause_flow_run(wait_for_input=NextStepInput)
    print(f"Received input: {res}")

    if res.next == "step2":
        payload = await step2(payload)
    elif res.next == "step3":
        payload = await step3(payload)
    else:
        print("No next step specified, ending flow.")
        return "No next step specified"



    # suspend_flow_run()
    res = await apause_flow_run()
    print(res)

    output_path = await step4(payload)
    print(f"Saved payload to {output_path}")
    return output_path


if __name__ == "__main__":
    map_pipeline_flow()