task_executor.py

#!/usr/bin/env python3
"""
Multi-Task Configuration Parser & Batch Executor

Parses a YAML/JSON task configuration and executes multiple independent tasks
in parallel with timeout enforcement, retry logic, structured logging,
graceful shutdown (SIGINT/SIGTERM), and DAG dependency ordering.

Usage:
    python task_executor.py config.yaml [--output results.json] [--log-level DEBUG]

Config format (YAML or JSON):
    version: "1.0"
    max_parallelism: 3
    log_level: "INFO"
    tasks:
      - id: task1
        command: "echo hello"
        timeout: 10        # seconds
        retries: 2
        retry_delay: 1.0   # seconds, exponential backoff base
        depends_on: []

Output:
    Structured JSON with status, stdout, stderr, exit_code, duration per task.
"""

from __future__ import annotations

import argparse
import concurrent.futures
import json
import logging
import os
import signal
import subprocess
import sys
import time
import uuid
from collections import defaultdict, deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

# Optional YAML support
try:
    import yaml
    HAS_YAML = True
except ImportError:
    HAS_YAML = False

# ---------------------------------------------------------------------------
# Globals / shutdown flag
# ---------------------------------------------------------------------------
_shutdown_requested = False
_executor_ref: concurrent.futures.ThreadPoolExecutor | None = None


def _signal_handler(signum: int, frame: Any) -> None:
    global _shutdown_requested
    _shutdown_requested = True
    logging.warning("Shutdown signal received (%s). Draining active tasks...", signum)
    if _executor_ref:
        _executor_ref.shutdown(wait=False, cancel_futures=True)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def load_config(path: str) -> dict[str, Any]:
    """Load task configuration from YAML or JSON file."""
    raw = Path(path).read_text(encoding="utf-8")
    if path.lower().endswith((".yaml", ".yml")):
        if not HAS_YAML:
            raise RuntimeError("PyYAML not installed. Install with: pip install pyyaml")
        return yaml.safe_load(raw)
    return json.loads(raw)


def topological_sort(tasks: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Return tasks in topological order respecting `depends_on` edges.

    Raises ValueError if a cycle is detected.
    """
    by_id = {t["id"]: t for t in tasks}
    in_degree: dict[str, int] = {t["id"]: 0 for t in tasks}
    dependents: dict[str, list[str]] = defaultdict(list)

    for task in tasks:
        for dep in task.get("depends_on") or []:
            if dep not in by_id:
                raise ValueError(f"Task '{task['id']}' depends on unknown task '{dep}'")
            in_degree[task["id"]] += 1
            dependents[dep].append(task["id"])

    queue: deque[str] = deque(tid for tid, deg in in_degree.items() if deg == 0)
    ordered: list[dict[str, Any]] = []

    while queue:
        tid = queue.popleft()
        ordered.append(by_id[tid])
        for dependent in dependents[tid]:
            in_degree[dependent] -= 1
            if in_degree[dependent] == 0:
                queue.append(dependent)

    if len(ordered) != len(tasks):
        cycle_ids = [tid for tid, deg in in_degree.items() if deg > 0]
        raise ValueError(f"Cycle detected in task dependencies: {cycle_ids}")

    return ordered


# ---------------------------------------------------------------------------
# Task runner
# ---------------------------------------------------------------------------

def run_single_task(
    task: dict[str, Any],
    logger: logging.Logger,
) -> dict[str, Any]:
    """Execute a single task with timeout + exponential-backoff retry."""
    task_id = task["id"]
    command = task["command"]
    timeout = float(task.get("timeout", 60))
    max_retries = int(task.get("retries", 0))
    retry_delay = float(task.get("retry_delay", 1.0))

    result: dict[str, Any] = {
        "id": task_id,
        "command": command,
        "status": "pending",
        "exit_code": None,
        "stdout": "",
        "stderr": "",
        "duration_seconds": 0.0,
        "retries_used": 0,
        "started_utc": None,
        "completed_utc": None,
        "error": None,
    }

    for attempt in range(max_retries + 1):
        if _shutdown_requested:
            result["status"] = "skipped"
            result["error"] = "Shutdown requested before task could run"
            return result

        if attempt > 0:
            backoff = retry_delay * (2 ** (attempt - 1))
            logger.info("[%s] Retry %d/%d after %.1fs backoff", task_id, attempt, max_retries, backoff)
            time.sleep(backoff)

        result["started_utc"] = utc_now()
        t0 = time.monotonic()

        try:
            proc = subprocess.run(
                command,
                shell=True,
                capture_output=True,
                text=True,
                timeout=timeout,
            )
            elapsed = time.monotonic() - t0
            result.update({
                "exit_code": proc.returncode,
                "stdout": proc.stdout,
                "stderr": proc.stderr,
                "duration_seconds": round(elapsed, 3),
                "retries_used": attempt,
                "completed_utc": utc_now(),
            })
            if proc.returncode == 0:
                result["status"] = "success"
                logger.info("[%s] Completed in %.3fs (exit 0)", task_id, elapsed)
            else:
                result["status"] = "failed"
                logger.warning("[%s] Failed exit=%d after %.3fs", task_id, proc.returncode, elapsed)
                if attempt < max_retries:
                    continue
            break

        except subprocess.TimeoutExpired:
            elapsed = time.monotonic() - t0
            result.update({
                "status": "timedout",
                "duration_seconds": round(elapsed, 3),
                "retries_used": attempt,
                "completed_utc": utc_now(),
                "error": f"Timed out after {timeout}s",
            })
            logger.error("[%s] Timed out after %.1fs", task_id, timeout)
            if attempt < max_retries:
                continue
            break

        except Exception as exc:
            elapsed = time.monotonic() - t0
            result.update({
                "status": "error",
                "duration_seconds": round(elapsed, 3),
                "retries_used": attempt,
                "completed_utc": utc_now(),
                "error": str(exc),
            })
            logger.error("[%s] Unexpected error: %s", task_id, exc)
            if attempt < max_retries:
                continue
            break

    return result


# ---------------------------------------------------------------------------
# Batch executor
# ---------------------------------------------------------------------------

def execute_batch(
    config: dict[str, Any],
    logger: logging.Logger,
) -> dict[str, Any]:
    """Execute all tasks in config respecting parallelism and dependency order."""
    global _executor_ref

    max_parallelism = int(config.get("max_parallelism", 4))
    tasks = config.get("tasks", [])

    if not tasks:
        raise ValueError("No tasks defined in configuration.")

    # Topological sort to respect dependencies
    ordered_tasks = topological_sort(tasks)

    run_id = str(uuid.uuid4())
    started_utc = utc_now()
    t_start = time.monotonic()

    results: dict[str, dict[str, Any]] = {}
    task_futures: dict[str, concurrent.futures.Future] = {}  # type: ignore[type-arg]

    logger.info("Starting batch run_id=%s tasks=%d parallelism=%d", run_id, len(ordered_tasks), max_parallelism)

    with concurrent.futures.ThreadPoolExecutor(max_workers=max_parallelism) as executor:
        _executor_ref = executor
        pending = list(ordered_tasks)
        submitted: set[str] = set()

        while pending or task_futures:
            if _shutdown_requested:
                logger.warning("Shutdown: cancelling %d pending tasks", len(pending))
                for future in task_futures.values():
                    future.cancel()
                break

            # Submit tasks whose dependencies are satisfied
            for task in list(pending):
                if task["id"] in submitted:
                    continue
                deps = task.get("depends_on") or []
                # Check all deps are done (success or skipped — not pending or failed)
                deps_ok = all(
                    results.get(dep, {}).get("status") in ("success",)
                    for dep in deps
                )
                deps_blocked = any(
                    results.get(dep, {}).get("status") in ("failed", "timedout", "error", "skipped")
                    for dep in deps
                )

                if not deps and task["id"] not in submitted:
                    # No deps — submit immediately
                    future = executor.submit(run_single_task, task, logger)
                    task_futures[task["id"]] = future
                    submitted.add(task["id"])
                    pending.remove(task)
                elif deps_ok:
                    future = executor.submit(run_single_task, task, logger)
                    task_futures[task["id"]] = future
                    submitted.add(task["id"])
                    pending.remove(task)
                elif deps_blocked:
                    # Dependency failed — skip this task
                    logger.warning("[%s] Skipping — dependency failed", task["id"])
                    results[task["id"]] = {
                        "id": task["id"],
                        "command": task.get("command", ""),
                        "status": "skipped",
                        "exit_code": None,
                        "stdout": "",
                        "stderr": "",
                        "duration_seconds": 0.0,
                        "retries_used": 0,
                        "started_utc": None,
                        "completed_utc": None,
                        "error": "Dependency failed",
                    }
                    submitted.add(task["id"])
                    pending.remove(task)

            # Collect completed futures
            done_ids = [tid for tid, fut in task_futures.items() if fut.done()]
            for tid in done_ids:
                future = task_futures.pop(tid)
                try:
                    results[tid] = future.result()
                except Exception as exc:
                    results[tid] = {
                        "id": tid, "status": "error", "error": str(exc),
                        "exit_code": None, "stdout": "", "stderr": "",
                        "duration_seconds": 0.0, "retries_used": 0,
                        "started_utc": None, "completed_utc": utc_now(),
                    }

            if task_futures or pending:
                time.sleep(0.05)

        _executor_ref = None

    total_elapsed = time.monotonic() - t_start
    all_results = [results.get(t["id"], {"id": t["id"], "status": "unknown"}) for t in ordered_tasks]

    summary = {
        "total": len(ordered_tasks),
        "success": sum(1 for r in all_results if r.get("status") == "success"),
        "failed": sum(1 for r in all_results if r.get("status") in ("failed", "timedout", "error")),
        "timedout": sum(1 for r in all_results if r.get("status") == "timedout"),
        "skipped": sum(1 for r in all_results if r.get("status") == "skipped"),
    }

    logger.info(
        "Batch complete in %.2fs — success=%d failed=%d timedout=%d skipped=%d",
        total_elapsed, summary["success"], summary["failed"], summary["timedout"], summary["skipped"]
    )

    return {
        "run_id": run_id,
        "started_utc": started_utc,
        "completed_utc": utc_now(),
        "total_duration_seconds": round(total_elapsed, 3),
        "config_max_parallelism": max_parallelism,
        "tasks": all_results,
        "summary": summary,
    }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Multi-Task Configuration Parser & Batch Executor",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("config", help="Path to YAML or JSON task config file")
    parser.add_argument(
        "--output", "-o", default="results.json",
        help="Output path for JSON results (default: results.json)"
    )
    parser.add_argument(
        "--log-level", default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Logging verbosity (default: INFO)"
    )
    parser.add_argument(
        "--dry-run", action="store_true",
        help="Parse config and show task plan without executing"
    )
    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s %(levelname)-8s %(message)s",
        datefmt="%Y-%m-%dT%H:%M:%SZ",
    )
    logger = logging.getLogger("task_executor")

    # Register signal handlers for graceful shutdown
    signal.signal(signal.SIGINT, _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    try:
        config = load_config(args.config)
    except Exception as exc:
        logger.error("Failed to load config '%s': %s", args.config, exc)
        return 1

    tasks = config.get("tasks", [])
    logger.info("Loaded %d task(s) from '%s'", len(tasks), args.config)

    if args.dry_run:
        try:
            ordered = topological_sort(tasks)
        except ValueError as exc:
            logger.error("Config validation error: %s", exc)
            return 1
        logger.info("Dry-run: task execution order:")
        for t in ordered:
            logger.info("  [%s] cmd=%r timeout=%s retries=%s deps=%s",
                        t["id"], t.get("command", ""), t.get("timeout", 60),
                        t.get("retries", 0), t.get("depends_on", []))
        return 0

    try:
        report = execute_batch(config, logger)
    except Exception as exc:
        logger.error("Batch execution error: %s", exc)
        return 1

    out_path = Path(args.output)
    out_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    logger.info("Results written to '%s'", out_path)

    failed = report["summary"]["failed"]
    return 0 if failed == 0 else 1


if __name__ == "__main__":
    raise SystemExit(main())