SupplyAgent/supply_infra/pipeline/orchestrator.py

from __future__ import annotations

from datetime import timedelta

from sqlalchemy.orm import Session

from supply_infra.db.repositories.pipeline_run_repo import PipelineRunRepository
from supply_infra.db.repositories.pipeline_step_run_repo import PipelineStepRunRepository
from supply_infra.pipeline.contracts import StepResult
from supply_infra.pipeline.dates import china_now
from supply_infra.pipeline.enums import RunStatus, StepStatus
from supply_infra.pipeline.gates import evaluate_step_gate

_TRANSIENT_ERRORS = {
    "step_timed_out",
    "step_launch_failed",
    "db_pool_exhausted",
}


def complete_step(
    session: Session,
    *,
    step_run_id: str,
    owner: str,
    result: StepResult,
) -> str:
    now = china_now()
    step_repo = PipelineStepRunRepository(session)
    run_repo = PipelineRunRepository(session)
    step = step_repo.get(step_run_id, for_update=True)
    if step is None:
        raise ValueError(f"step_run_id not found: {step_run_id}")
    if step.status != StepStatus.RUNNING.value or step.lease_owner != owner:
        raise RuntimeError(f"Step lease lost: {step_run_id}")
    run = run_repo.get(step.run_id, for_update=True)
    if run is None:
        raise ValueError(f"pipeline run not found: {step.run_id}")
    if run.status in {
        RunStatus.CANCELLING.value,
        RunStatus.CANCELLED.value,
    }:
        step_repo.finish(
            step,
            status=StepStatus.CANCELLED.value,
            now=now,
            exit_code=result.exit_code,
            result=result.payload,
            error_code="cancelled",
            error_message="Step result discarded after pipeline cancellation",
            log_uri=result.log_uri,
        )
        if run.status == RunStatus.CANCELLING.value:
            run_repo.finish(run, status=RunStatus.CANCELLED.value, now=now)
        return RunStatus.CANCELLED.value
    if run.status == RunStatus.DEADLINE_EXCEEDED.value:
        step_repo.finish(
            step,
            status=StepStatus.TIMED_OUT.value,
            now=now,
            exit_code=result.exit_code,
            result=result.payload,
            error_code="deadline_exceeded",
            error_message="Step result discarded after pipeline deadline",
            log_uri=result.log_uri,
        )
        return RunStatus.DEADLINE_EXCEEDED.value

    gate = evaluate_step_gate(step.step_key, result.payload) if result.success else None
    passed = bool(result.success and gate and gate.passed)
    error_code = result.error_code
    error_message = result.error_message
    if result.success and gate is not None and not gate.passed:
        error_code = gate.error_code
        error_message = gate.error_message

    if passed:
        step_repo.finish(
            step,
            status=StepStatus.SUCCEEDED.value,
            now=now,
            exit_code=result.exit_code,
            result=result.payload,
            error_code=None,
            error_message=None,
            log_uri=result.log_uri,
        )
        latest = step_repo.latest_attempts(run.run_id)
        incomplete = [
            item
            for item in sorted(latest.values(), key=lambda value: value.step_order)
            if item.status != StepStatus.SUCCEEDED.value
        ]
        if not incomplete:
            run_repo.finish(
                run,
                status=RunStatus.SUCCEEDED.value,
                now=now,
                summary={"steps": len(latest), "dry_run": run.dry_run},
            )
            return RunStatus.SUCCEEDED.value
        next_step = incomplete[0]
        if next_step.status == StepStatus.PENDING.value:
            next_step.status = StepStatus.READY.value
            next_step.error_code = None
            next_step.error_message = None
        run.status = RunStatus.QUEUED.value
        run.current_step = next_step.step_key
        run.heartbeat_at = now
        run.lease_owner = None
        run.lease_until = None
        session.flush()
        return RunStatus.QUEUED.value

    failed_status = (
        StepStatus.TIMED_OUT.value if result.timed_out else StepStatus.FAILED.value
    )
    step_repo.finish(
        step,
        status=failed_status,
        now=now,
        exit_code=result.exit_code,
        result=result.payload,
        error_code=error_code or "step_failed",
        error_message=error_message,
        log_uri=result.log_uri,
    )
    should_retry = (
        step.retryable
        and step.attempt < step.max_attempts
        and (error_code or "") in _TRANSIENT_ERRORS
    )
    if should_retry:
        retry_at = now + timedelta(seconds=min(300, 30 * (2 ** (step.attempt - 1))))
        step_repo.create_retry_attempt(
            step,
            status=StepStatus.RETRY_WAIT.value,
            next_retry_at=retry_at,
        )
        run.status = RunStatus.QUEUED.value
        run.current_step = step.step_key
        run.heartbeat_at = now
        run.lease_owner = None
        run.lease_until = None
        session.flush()
        return StepStatus.RETRY_WAIT.value

    reason = error_message or error_code or "critical step failed"
    step_repo.block_unfinished_downstream(
        run_id=run.run_id,
        after_order=step.step_order,
        reason=reason,
    )
    run_repo.finish(
        run,
        status=RunStatus.FAILED.value,
        now=now,
        error_code=error_code or "critical_step_failed",
        error_message=reason,
        summary={"failed_step": step.step_key, "attempt": step.attempt},
    )
    return RunStatus.FAILED.value