from __future__ import annotations import logging import os import signal import socket import threading from datetime import timedelta from zoneinfo import ZoneInfo from supply_infra.config import InfraSettings, get_infra_settings from supply_infra.db.repositories.pipeline_run_repo import PipelineRunRepository from supply_infra.db.repositories.pipeline_step_run_repo import PipelineStepRunRepository from supply_infra.db.session import dispose_engine, get_session from supply_infra.pipeline.dates import CHINA_TIMEZONE, china_now from supply_infra.pipeline.enums import RunStatus, StepStatus from supply_infra.pipeline.health import ( run_alert_level, touch_reconciler_heartbeat, ) from supply_infra.pipeline.run_service import PIPELINE_KEY, submit_pipeline_run logger = logging.getLogger(__name__) def reconcile_once() -> dict[str, int]: settings = get_infra_settings() now = china_now() stats = { "retries_promoted": 0, "expired_steps": 0, "deadline_exceeded": 0, "previous_run_unblocked": 0, "alerts_emitted": 0, "missed_runs_submitted": 0, } with get_session() as session: run_repo = PipelineRunRepository(session) step_repo = PipelineStepRunRepository(session) stats["retries_promoted"] = step_repo.promote_due_retries(now) for step in step_repo.list_expired_running(now): run = run_repo.get(step.run_id, for_update=True) if run is not None and run.status in { RunStatus.CANCELLING.value, RunStatus.CANCELLED.value, }: step.status = StepStatus.CANCELLED.value step.finished_at = now step.lease_owner = None step.lease_until = None step.error_code = "cancelled" step.error_message = "Cancellation completed after worker lease expired" if run.status == RunStatus.CANCELLING.value: run_repo.finish( run, status=RunStatus.CANCELLED.value, now=now, ) stats["expired_steps"] += 1 continue if run is not None and run.status == RunStatus.DEADLINE_EXCEEDED.value: step.status = StepStatus.TIMED_OUT.value step.finished_at = now step.lease_owner = None step.lease_until = None step.error_code = "deadline_exceeded" step.error_message = "Worker lease expired after pipeline deadline" stats["expired_steps"] += 1 continue step.status = StepStatus.TIMED_OUT.value step.finished_at = now step.lease_owner = None step.lease_until = None step.error_code = "lease_expired" step.error_message = "Worker heartbeat lease expired" if step.retryable and step.attempt < step.max_attempts: step_repo.create_retry_attempt(step, status=StepStatus.READY.value) if run is not None: run.status = RunStatus.QUEUED.value run.current_step = step.step_key run.lease_owner = None run.lease_until = None else: step_repo.block_unfinished_downstream( run_id=step.run_id, after_order=step.step_order, reason="Worker heartbeat lease expired", ) if run is not None: run_repo.finish( run, status=RunStatus.FAILED.value, now=now, error_code="lease_expired", error_message=f"Lease expired at step {step.step_key}", ) stats["expired_steps"] += 1 for run in run_repo.list_past_deadline(now): for step in step_repo.list_for_run(run.run_id): if step.status == StepStatus.RUNNING.value: step.error_code = "deadline_exceeded" step.error_message = ( "Pipeline exceeded next scheduled start; waiting for " "the worker or lease expiry" ) continue if step.status in { StepStatus.PENDING.value, StepStatus.READY.value, StepStatus.RETRY_WAIT.value, StepStatus.BLOCKED.value, }: step.status = StepStatus.BLOCKED.value step.error_code = "deadline_exceeded" step.error_message = "Pipeline exceeded next scheduled start" run_repo.finish( run, status=RunStatus.DEADLINE_EXCEEDED.value, now=now, error_code="deadline_exceeded", error_message="Pipeline exceeded next scheduled start", ) stats["deadline_exceeded"] += 1 for run in run_repo.list_blocked_previous(): blocked_by = (run.summary_json or {}).get("blocked_by_run_id") previous = run_repo.get(str(blocked_by)) if blocked_by else None if previous is not None and previous.status in { RunStatus.QUEUED.value, RunStatus.RUNNING.value, RunStatus.CANCELLING.value, RunStatus.BLOCKED_PREVIOUS_RUN.value, }: continue if previous is not None and step_repo.has_running_for_run(previous.run_id): continue latest = step_repo.latest_attempts(run.run_id) first = min(latest.values(), key=lambda item: item.step_order, default=None) if first is None: continue first.status = StepStatus.READY.value first.error_code = None first.error_message = None for step in latest.values(): if step.step_run_id != first.step_run_id: step.status = StepStatus.PENDING.value step.error_code = None step.error_message = None run.status = RunStatus.QUEUED.value run.summary_json = None stats["previous_run_unblocked"] += 1 for run in run_repo.list_nonterminal(): alert_level = run_alert_level( status=run.status, started_at=run.started_at, created_at=run.created_at, deadline_at=run.deadline_at, now=now, settings=settings, ) if alert_level is not None: logger.warning( "pipeline_alert level=%s run_id=%s biz_dt=%s status=%s deadline=%s", alert_level, run.run_id, run.biz_dt, run.status, run.deadline_at.isoformat() if run.deadline_at else None, ) stats["alerts_emitted"] += 1 local_now = now.replace(tzinfo=CHINA_TIMEZONE).astimezone( ZoneInfo(settings.scheduler_timezone) ) missed_cutoff = local_now.replace( hour=settings.scheduler_cron_hour, minute=settings.scheduler_cron_minute, second=0, microsecond=0, ) + timedelta(seconds=settings.pipeline_missed_run_grace_seconds) if settings.scheduler_enabled and local_now >= missed_cutoff: biz_dt = local_now.strftime("%Y%m%d") with get_session() as session: existing = PipelineRunRepository(session).get_for_business_day( pipeline_key=PIPELINE_KEY, biz_dt=biz_dt, ) if existing is None: submission = submit_pipeline_run( biz_dt=biz_dt, trigger_type="reconcile", trigger_source="missed_run_reconciler", trigger_reason="Scheduled batch was absent after the grace window", settings=settings, ) if submission.created: logger.error( "pipeline_alert level=critical type=missed_run_recovered " "run_id=%s biz_dt=%s", submission.run_id, biz_dt, ) stats["missed_runs_submitted"] = 1 return stats class PipelineReconciler: def __init__(self, settings: InfraSettings | None = None) -> None: self.settings = settings or get_infra_settings() self.owner = f"{socket.gethostname()}:{os.getpid()}" self._stop = threading.Event() def stop(self, *_args: object) -> None: self._stop.set() def run_forever(self) -> None: signal.signal(signal.SIGTERM, self.stop) signal.signal(signal.SIGINT, self.stop) logger.info("Pipeline reconciler started") try: while not self._stop.is_set(): try: result = reconcile_once() touch_reconciler_heartbeat( owner=self.owner, settings=self.settings, ) logger.info("Pipeline reconcile result: %s", result) except Exception: logger.exception("Pipeline reconcile failed") self._stop.wait(self.settings.pipeline_reconcile_seconds) finally: dispose_engine() logger.info("Pipeline reconciler stopped")