125 lines
4.0 KiB
Python
Executable File
125 lines
4.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Simulate placement for AI-ish homelab workloads."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
from pathlib import Path
|
|
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
|
SCHED_DIR = REPO_ROOT / "infra" / "ai-scheduler"
|
|
NODES_FILE = SCHED_DIR / "nodes.tsv"
|
|
WORKLOADS_FILE = SCHED_DIR / "workloads.tsv"
|
|
|
|
|
|
def read_tsv(path: Path) -> list[dict[str, str]]:
|
|
with path.open(encoding="utf-8", newline="") as handle:
|
|
return list(csv.DictReader(handle, delimiter="\t"))
|
|
|
|
|
|
def as_bool(value: str) -> bool:
|
|
return value.strip().lower() in {"true", "yes", "1"}
|
|
|
|
|
|
def score_node(workload: dict[str, str], node: dict[str, str]) -> tuple[int, list[str], list[str]]:
|
|
score = 100
|
|
reasons: list[str] = []
|
|
blockers: list[str] = []
|
|
cpu_needed = int(workload["cpu_m"])
|
|
mem_needed = int(workload["memory_mib"])
|
|
cpu_have = int(node["cpu_m"])
|
|
mem_have = int(node["memory_mib"])
|
|
|
|
if cpu_have < cpu_needed:
|
|
blockers.append(f"cpu {cpu_have}m < {cpu_needed}m")
|
|
else:
|
|
reasons.append("cpu fits")
|
|
score += min((cpu_have - cpu_needed) // 250, 10)
|
|
if mem_have < mem_needed:
|
|
blockers.append(f"memory {mem_have}Mi < {mem_needed}Mi")
|
|
else:
|
|
reasons.append("memory fits")
|
|
score += min((mem_have - mem_needed) // 512, 10)
|
|
if as_bool(workload["requires_tailnet"]) and not as_bool(node["tailnet"]):
|
|
blockers.append("tailnet required")
|
|
elif as_bool(node["tailnet"]):
|
|
reasons.append("tailnet available")
|
|
score += 8
|
|
if workload["prefer_arch"] != "any" and workload["prefer_arch"] != node["arch"]:
|
|
score -= 15
|
|
reasons.append(f"non-preferred arch {node['arch']}")
|
|
else:
|
|
reasons.append(f"arch {node['arch']} ok")
|
|
if as_bool(node["control_plane"]) and not as_bool(workload["allow_control_plane"]):
|
|
score -= 30
|
|
reasons.append("control-plane penalty")
|
|
if node["power_class"] == "low":
|
|
score += 5
|
|
reasons.append("low-power node")
|
|
return score, reasons, blockers
|
|
|
|
|
|
def recommend(workload_name: str) -> int:
|
|
workloads = {row["workload"]: row for row in read_tsv(WORKLOADS_FILE)}
|
|
nodes = read_tsv(NODES_FILE)
|
|
if workload_name not in workloads:
|
|
print(f"Unknown workload: {workload_name}")
|
|
print("Known workloads:")
|
|
for name in workloads:
|
|
print(f" {name}")
|
|
return 2
|
|
workload = workloads[workload_name]
|
|
rows = []
|
|
for node in nodes:
|
|
score, reasons, blockers = score_node(workload, node)
|
|
rows.append((score, node["node"], reasons, blockers))
|
|
rows.sort(reverse=True)
|
|
|
|
print("AI Workload Scheduler")
|
|
print("=====================")
|
|
print(f"workload: {workload_name}")
|
|
print()
|
|
for score, node_name, reasons, blockers in rows:
|
|
status = "blocked" if blockers else "candidate"
|
|
print(f"{status:9} {node_name:18} score={score}")
|
|
if blockers:
|
|
print(f" blockers: {', '.join(blockers)}")
|
|
print(f" reasons: {', '.join(reasons)}")
|
|
best = next((row for row in rows if not row[3]), None)
|
|
if best:
|
|
print()
|
|
print(f"Recommendation: {best[1]}")
|
|
return 0
|
|
print()
|
|
print("Recommendation: add capacity or relax workload requirements")
|
|
return 1
|
|
|
|
|
|
def simulate() -> int:
|
|
failures = 0
|
|
for workload in read_tsv(WORKLOADS_FILE):
|
|
print()
|
|
rc = recommend(workload["workload"])
|
|
failures += 1 if rc else 0
|
|
return 1 if failures else 0
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
subparsers = parser.add_subparsers(dest="command", required=True)
|
|
subparsers.add_parser("simulate")
|
|
rec = subparsers.add_parser("recommend")
|
|
rec.add_argument("workload")
|
|
args = parser.parse_args()
|
|
if args.command == "simulate":
|
|
return simulate()
|
|
if args.command == "recommend":
|
|
return recommend(args.workload)
|
|
return 2
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|