pipeline_runner.py

#!/usr/bin/env python3
"""
Dependency-Aware Pipeline Runner
Executes jobs based on their dependencies using topological sorting
"""

import time
import yaml
import subprocess
from typing import Dict, List
from enum import Enum
from collections import defaultdict, deque


class JobStatus(Enum):
    PENDING = "pending"
    RUNNING = "running"
    SUCCESS = "success"
    FAILED = "failed"
    SKIPPED = "skipped"


class Job:
    """Represents a single job in the pipeline"""

    def __init__(self, name: str, command: str, dependencies: List[str] = None):
        self.name = name
        self.command = command
        self.dependencies = dependencies or []
        self.status = JobStatus.PENDING
        self.start_time = None
        self.end_time = None
        self.error_message = None

    def execute(self) -> bool:
        """
        Execute the job command.

        Returns:
        bool: True if successful, False otherwise
        """
        self.status = JobStatus.RUNNING
        self.start_time = time.time()

        print(f"\n[RUNNING] {self.name}")
        print(f"Command: {self.command}")

        try:
            result = subprocess.run(
                self.command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=30
            )

            if result.stdout and result.stdout.strip():
                print(result.stdout.strip())

            if result.returncode == 0:
                self.status = JobStatus.SUCCESS
            else:
                self.status = JobStatus.FAILED
                stderr_message = result.stderr.strip() if result.stderr else ""
                self.error_message = stderr_message or f"Command failed with exit code {result.returncode}"

            self.end_time = time.time()
            return self.status == JobStatus.SUCCESS

        except subprocess.TimeoutExpired:
            self.status = JobStatus.FAILED
            self.error_message = "Command timed out after 30 seconds"
            self.end_time = time.time()
            return False

        except Exception as e:
            self.status = JobStatus.FAILED
            self.error_message = str(e)
            self.end_time = time.time()
            return False

    def __repr__(self):
        return f"Job({self.name}, status={self.status.value})"


class PipelineRunner:
    """Manages and executes jobs based on dependencies"""

    def __init__(self, config_file: str):
        self.jobs: Dict[str, Job] = {}
        self.dependency_graph: Dict[str, List[str]] = defaultdict(list)
        self.reverse_graph: Dict[str, List[str]] = defaultdict(list)
        self.load_config(config_file)

    def load_config(self, config_file: str):
        """
        Load pipeline configuration from YAML file.

        - Read YAML file
        - Create Job objects
        - Build dependency_graph (job -> dependents)
        - Build reverse_graph (job -> dependencies)
        """
        with open(config_file, "r", encoding="utf-8") as file_obj:
            config = yaml.safe_load(file_obj)

        if not config or "jobs" not in config:
            raise ValueError("Invalid configuration: missing 'jobs' section")

        for job_data in config["jobs"]:
            name = job_data.get("name")
            command = job_data.get("command")
            dependencies = job_data.get("dependencies", [])

            if not name:
                raise ValueError("Job configuration missing 'name'")
            if not command:
                raise ValueError(f"Job '{name}' is missing 'command'")
            if name in self.jobs:
                raise ValueError(f"Duplicate job name found: {name}")

            self.jobs[name] = Job(name, command, dependencies)
            self.reverse_graph[name] = list(dependencies)

        for job_name in self.jobs:
            self.dependency_graph[job_name] = self.dependency_graph.get(job_name, [])
            self.reverse_graph[job_name] = self.reverse_graph.get(job_name, [])

        for job_name, job in self.jobs.items():
            for dependency in job.dependencies:
                if dependency not in self.jobs:
                    raise ValueError(
                        f"Job '{job_name}' depends on undefined job '{dependency}'"
                    )
                self.dependency_graph[dependency].append(job_name)

    def validate_dag(self) -> bool:
        """
        Validate that the dependency graph is a DAG (no cycles).

        Returns:
        bool: True if valid DAG, False if cycles detected
        """
        WHITE, GRAY, BLACK = 0, 1, 2
        color = {job: WHITE for job in self.jobs}

        def has_cycle(node: str) -> bool:
            color[node] = GRAY

            for neighbor in self.dependency_graph[node]:
                if color[neighbor] == GRAY:
                    return True
                if color[neighbor] == WHITE and has_cycle(neighbor):
                    return True

            color[node] = BLACK
            return False

        for job in self.jobs:
            if color[job] == WHITE:
                if has_cycle(job):
                    return False

        return True

    def topological_sort(self) -> List[str]:
        """
        Perform topological sort using Kahn's algorithm.

        Returns:
        List[str]: Ordered list of job names
        """
        in_degree = {job: 0 for job in self.jobs}

        for job_name in self.jobs:
            in_degree[job_name] = len(self.reverse_graph[job_name])

        queue = deque(sorted([job for job, degree in in_degree.items() if degree == 0]))
        sorted_order = []

        while queue:
            current = queue.popleft()
            sorted_order.append(current)

            for dependent in sorted(self.dependency_graph[current]):
                in_degree[dependent] -= 1
                if in_degree[dependent] == 0:
                    queue.append(dependent)

        if len(sorted_order) != len(self.jobs):
            raise ValueError("Topological sort failed: graph may contain a cycle")

        return sorted_order

    def can_execute(self, job_name: str) -> bool:
        """
        Check if a job's dependencies are satisfied.

        - Check all dependencies have SUCCESS status
        - Return True only if all dependencies succeeded
        """
        dependencies = self.reverse_graph[job_name]

        for dependency in dependencies:
            if self.jobs[dependency].status != JobStatus.SUCCESS:
                return False

        return True

    def mark_downstream_skipped(self, job_name: str):
        """
        Mark all downstream jobs as skipped when a job fails.

        - Use BFS to find all dependent jobs
        - Set their status to SKIPPED
        """
        queue = deque(self.dependency_graph[job_name])
        visited = set()

        while queue:
            current = queue.popleft()

            if current in visited:
                continue

            visited.add(current)

            if self.jobs[current].status == JobStatus.PENDING:
                self.jobs[current].status = JobStatus.SKIPPED
                if not self.jobs[current].error_message:
                    self.jobs[current].error_message = (
                        f"Skipped because dependency '{job_name}' failed"
                    )

            for downstream in self.dependency_graph[current]:
                if downstream not in visited:
                    queue.append(downstream)

    def execute_pipeline(self) -> Dict[str, JobStatus]:
        """
        Execute pipeline in dependency order.

        Returns:
        Dict[str, JobStatus]: Final status of all jobs
        """
        print("\n" + "=" * 60)
        print("STARTING PIPELINE EXECUTION")
        print("=" * 60)

        if not self.validate_dag():
            print("ERROR: Cycle detected in dependencies!")
            return {}

        execution_order = self.topological_sort()
        print(f"\nExecution order: {' -> '.join(execution_order)}")

        for job_name in execution_order:
            job = self.jobs[job_name]

            if job.status == JobStatus.SKIPPED:
                continue

            if not self.can_execute(job_name):
                if job.status == JobStatus.PENDING:
                    job.status = JobStatus.SKIPPED
                    failed_dependencies = [
                        dep for dep in self.reverse_graph[job_name]
                        if self.jobs[dep].status != JobStatus.SUCCESS
                    ]
                    job.error_message = (
                        f"Skipped because dependencies not satisfied: {', '.join(failed_dependencies)}"
                    )
                continue

            success = job.execute()

            if not success:
                self.mark_downstream_skipped(job_name)

        return {name: job.status for name, job in self.jobs.items()}

    def print_summary(self):
        """Print execution summary"""
        print("\n" + "=" * 60)
        print("PIPELINE EXECUTION SUMMARY")
        print("=" * 60)

        for job_name, job in self.jobs.items():
            duration = ""
            if job.start_time and job.end_time:
                duration = f" ({job.end_time - job.start_time:.2f}s)"

            status_symbol = {
                JobStatus.SUCCESS: "✓",
                JobStatus.FAILED: "✗",
                JobStatus.SKIPPED: "⊘",
                JobStatus.PENDING: "○",
                JobStatus.RUNNING: "▶"
            }.get(job.status, "?")

            print(f"{status_symbol} {job_name}: {job.status.value}{duration}")
            if job.error_message:
                print(f"  Error: {job.error_message}")

        print("=" * 60)


def main():
    """Main entry point"""
    import sys

    if len(sys.argv) < 2:
        print("Usage: python3 pipeline_runner.py <config_file>")
        sys.exit(1)

    config_file = sys.argv[1]

    runner = PipelineRunner(config_file)
    runner.execute_pipeline()
    runner.print_summary()


if __name__ == "__main__":
    main()