SupplyAgent/supply_infra/pipeline/dag.py

from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class StepDefinition:
    key: str
    order: int
    dependencies: tuple[str, ...]
    timeout_seconds: int
    max_attempts: int = 1
    retryable: bool = False
    critical: bool = True


_RAW_STEPS: tuple[tuple[str, int, bool], ...] = (
    ("global_tree_sync", 3600, True),
    ("demand_pool_source_sync", 3600, True),
    ("demand_classify", 21600, False),
    ("demand_belong_rel_sync", 3600, True),
    ("real_metrics_sync", 3600, True),
    ("popularity_stats", 3600, True),
    ("category_tree_weight", 3600, True),
    ("source_video_sync", 7200, True),
    ("demand_grade", 21600, False),
    ("demand_expand", 21600, False),
    ("video_discovery", 21600, False),
    ("aigc_write_record", 1800, True),
)


def _build_steps() -> tuple[StepDefinition, ...]:
    result: list[StepDefinition] = []
    previous: str | None = None
    for order, (key, timeout_seconds, retryable) in enumerate(_RAW_STEPS, start=1):
        result.append(
            StepDefinition(
                key=key,
                order=order,
                dependencies=(previous,) if previous else (),
                timeout_seconds=timeout_seconds,
                max_attempts=2 if retryable else 1,
                retryable=retryable,
            )
        )
        previous = key
    return tuple(result)


PIPELINE_STEPS = _build_steps()
STEP_BY_KEY = {step.key: step for step in PIPELINE_STEPS}


def validate_dag() -> None:
    seen: set[str] = set()
    for step in PIPELINE_STEPS:
        if step.key in seen:
            raise ValueError(f"Duplicate step key: {step.key}")
        missing = set(step.dependencies) - seen
        if missing:
            raise ValueError(f"Step {step.key} has unresolved dependencies: {sorted(missing)}")
        seen.add(step.key)


validate_dag()