timmy-home/scripts/dynamic_dispatch_optimizer.py

#!/usr/bin/env python3
"""Dynamic dispatch optimizer for fleet-wide coordination.

Refs: timmy-home #552

Takes a fleet dispatch spec plus optional failover status and produces a
capacity-aware assignment plan. Safe by default: it prints the plan and only
writes an output file when explicitly requested.
"""

from __future__ import annotations

import argparse
import json
import os
from pathlib import Path
from typing import Any

TIMMY_HOME = Path(os.environ.get("TIMMY_HOME", Path.home() / ".timmy"))
STATUS_FILE = TIMMY_HOME / "failover_status.json"
SPEC_FILE = TIMMY_HOME / "fleet_dispatch.json"
OUTPUT_FILE = TIMMY_HOME / "dispatch_plan.json"


def load_json(path: Path, default: Any):
    if not path.exists():
        return default
    return json.loads(path.read_text())


def _host_status(host: dict[str, Any], failover_status: dict[str, Any]) -> str:
    if host.get("always_available"):
        return "ONLINE"
    fleet = failover_status.get("fleet") or {}
    return str(fleet.get(host["name"], "ONLINE")).upper()


def _lane_matches(host: dict[str, Any], lane: str) -> bool:
    host_lanes = set(host.get("lanes") or ["general"])
    if host.get("always_available", False):
        return True
    if lane == "general":
        return "general" in host_lanes
    return lane in host_lanes


def _choose_candidate(task: dict[str, Any], hosts: list[dict[str, Any]]):
    lane = task.get("lane", "general")
    preferred = task.get("preferred_hosts") or []

    preferred_map = {host["name"]: host for host in hosts}
    for host_name in preferred:
        host = preferred_map.get(host_name)
        if not host:
            continue
        if host["remaining_capacity"] <= 0:
            continue
        if _lane_matches(host, lane):
            return host

    matching = [host for host in hosts if host["remaining_capacity"] > 0 and _lane_matches(host, lane)]
    if matching:
        matching.sort(key=lambda host: (host["assigned_count"], -host["remaining_capacity"], host["name"]))
        return matching[0]

    fallbacks = [host for host in hosts if host["remaining_capacity"] > 0 and host.get("always_available")]
    if fallbacks:
        fallbacks.sort(key=lambda host: (host["assigned_count"], -host["remaining_capacity"], host["name"]))
        return fallbacks[0]

    return None


def generate_plan(spec: dict[str, Any], failover_status: dict[str, Any] | None = None) -> dict[str, Any]:
    failover_status = failover_status or {}
    raw_hosts = spec.get("hosts") or []
    tasks = list(spec.get("tasks") or [])

    online_hosts = []
    offline_hosts = []
    for host in raw_hosts:
        normalized = {
            "name": host["name"],
            "capacity": int(host.get("capacity", 1)),
            "remaining_capacity": int(host.get("capacity", 1)),
            "assigned_count": 0,
            "lanes": list(host.get("lanes") or ["general"]),
            "always_available": bool(host.get("always_available", False)),
            "status": _host_status(host, failover_status),
        }
        if normalized["status"] == "ONLINE":
            online_hosts.append(normalized)
        else:
            offline_hosts.append(normalized["name"])

    ordered_tasks = sorted(
        tasks,
        key=lambda item: (-int(item.get("priority", 0)), str(item.get("id", ""))),
    )

    assignments = []
    unassigned = []
    for task in ordered_tasks:
        candidate = _choose_candidate(task, online_hosts)
        if candidate is None:
            unassigned.append({
                "task_id": task.get("id"),
                "reason": f"no_online_host_for_lane:{task.get('lane', 'general')}",
            })
            continue

        candidate["remaining_capacity"] -= 1
        candidate["assigned_count"] += 1
        assignments.append({
            "task_id": task.get("id"),
            "host": candidate["name"],
            "lane": task.get("lane", "general"),
            "priority": int(task.get("priority", 0)),
        })

    return {
        "assignments": assignments,
        "offline_hosts": sorted(offline_hosts),
        "unassigned": unassigned,
    }


def write_plan(plan: dict[str, Any], output_path: Path):
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(plan, indent=2))


def parse_args():
    parser = argparse.ArgumentParser(description="Generate a fleet dispatch plan from host health and task demand.")
    parser.add_argument("--spec-file", type=Path, default=SPEC_FILE, help="JSON fleet spec with hosts[] and tasks[]")
    parser.add_argument("--status-file", type=Path, default=STATUS_FILE, help="Failover monitor JSON payload")
    parser.add_argument("--output", type=Path, default=OUTPUT_FILE, help="Output path for the generated plan")
    parser.add_argument("--write-output", action="store_true", help="Persist the generated plan to --output")
    parser.add_argument("--json", action="store_true", help="Print JSON only")
    return parser.parse_args()


def main():
    args = parse_args()
    spec = load_json(args.spec_file, {"hosts": [], "tasks": []})
    failover_status = load_json(args.status_file, {})
    plan = generate_plan(spec, failover_status)

    if args.write_output:
        write_plan(plan, args.output)

    if args.json:
        print(json.dumps(plan, indent=2))
        return

    print("--- Dynamic Dispatch Optimizer ---")
    print(f"Assignments: {len(plan['assignments'])}")
    if plan["offline_hosts"]:
        print("Offline hosts: " + ", ".join(plan["offline_hosts"]))
    for assignment in plan["assignments"]:
        print(f"- {assignment['task_id']} -> {assignment['host']} ({assignment['lane']}, p={assignment['priority']})")
    if plan["unassigned"]:
        print("Unassigned:")
        for item in plan["unassigned"]:
            print(f"- {item['task_id']}: {item['reason']}")
    if args.write_output:
        print(f"Wrote plan to {args.output}")


if __name__ == "__main__":
    main()