SupplyAgent/supply_infra/scheduler/jobs/run_supply_pipeline.py

"""Compatibility entry point for the durable supply pipeline."""
from __future__ import annotations

import argparse
import time
from typing import Any

from supply_infra.pipeline.run_service import get_pipeline_run, submit_pipeline_run
from supply_infra.scheduler.cli_result import run_cli

_TERMINAL = {
    "succeeded",
    "failed",
    "cancelled",
    "deadline_exceeded",
}


def run_supply_pipeline(
    biz_dt: str | None = None,
    *,
    wait: bool = False,
    wait_timeout_seconds: int = 86400,
    poll_seconds: float = 2.0,
) -> dict[str, Any]:
    """
    Submit a durable batch.

    This function no longer executes business steps directly. ``wait=True`` is
    retained for diagnostics and waits on MySQL state rather than an in-memory
    thread.
    """
    submission = submit_pipeline_run(
        biz_dt=biz_dt,
        trigger_type="cli",
        trigger_source="legacy_run_supply_pipeline",
    )
    accepted = submission.to_dict()
    if not wait:
        return accepted

    deadline = time.monotonic() + wait_timeout_seconds
    while time.monotonic() < deadline:
        payload = get_pipeline_run(submission.run_id)
        if payload is None:
            return {
                **accepted,
                "success": False,
                "error": "Persistent run disappeared",
            }
        if payload["status"] in _TERMINAL:
            return {
                **payload,
                "success": payload["status"] == "succeeded",
            }
        time.sleep(max(0.1, poll_seconds))
    return {
        **accepted,
        "success": False,
        "error": f"Timed out waiting for run after {wait_timeout_seconds}s",
    }


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Submit the durable supply pipeline")
    parser.add_argument("biz_dt", nargs="?")
    parser.add_argument("--wait", action="store_true")
    args = parser.parse_args()
    run_cli(
        lambda: run_supply_pipeline(args.biz_dt, wait=args.wait),
        label="run_supply_pipeline",
    )