scripts/claim3_dag.py

#!/usr/bin/env python3
"""Claim 3 - Dependency-aware DAG decomposition + 3-stage pipeline (Algorithm 1).

Implements the paper's single Algorithm 1 ("Dependency-Aware Subtask Scheduling
with Budget-Adaptive Routing"), CPU-only:

  Stage 1: (T,E) <- Decompose(Q; M_P); G <- ValidateAndRepair(T,E)   -> fallback to
           a linear CHAIN if the DAG is invalid (cycle, dangling edge, or > n_max
           nodes). Repair budget R_max = 2, planner cap n_max = 7.
  Stage 2: frontier = in-degree-0 subtasks; while frontier not empty, pop a ready
           subtask, predict utility u_hat, route to CLOUD iff u_hat > tau_t else EDGE,
           update C_used, push newly-unblocked subtasks (respects dependencies).
  Stage 3: aggregate sub-results in TOPOLOGICAL order.

NOTE: there is only ONE algorithm (Algorithm 1). The decomposition/validation is
Definition C.2 + a prompt-based XML planner (Fig 6), NOT a numbered "Algorithm 2".
The LLM planner quality needs an LLM; the scheduling/parallelism mechanism does not
and is what we reproduce here on synthetic DAGs.

Writes outputs/claim3.json and figs/claim3_dag.png.
"""
import json
import os
from collections import deque

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt

HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
OUT = os.path.join(HERE, "outputs")
FIG = os.path.join(HERE, "figs")
os.makedirs(OUT, exist_ok=True)
os.makedirs(FIG, exist_ok=True)

N_MAX = 7  # planner cap on number of subtasks
R_MAX = 2  # repair-attempt budget


def is_valid_dag(nodes, edges):
    """Valid iff: <= n_max nodes, all edge endpoints exist, and acyclic."""
    if len(nodes) == 0 or len(nodes) > N_MAX:
        return False, f"node count {len(nodes)} not in 1..{N_MAX}"
    nodeset = set(nodes)
    for a, b in edges:
        if a not in nodeset or b not in nodeset:
            return False, f"dangling edge {a}->{b}"
    # cycle check via Kahn
    indeg = {n: 0 for n in nodes}
    for a, b in edges:
        indeg[b] += 1
    q = deque([n for n in nodes if indeg[n] == 0])
    seen = 0
    indeg2 = dict(indeg)
    adj = {n: [] for n in nodes}
    for a, b in edges:
        adj[a].append(b)
    while q:
        x = q.popleft()
        seen += 1
        for y in adj[x]:
            indeg2[y] -= 1
            if indeg2[y] == 0:
                q.append(y)
    if seen != len(nodes):
        return False, "cycle detected"
    return True, "ok"


def validate_and_repair(nodes, edges):
    """Return (G_nodes, G_edges, mode, attempts). Fallback to a linear chain on
    failure after R_max repair attempts (Stage 1)."""
    attempts = 0
    ok, why = is_valid_dag(nodes, edges)
    cur_edges = list(edges)
    while not ok and attempts < R_MAX:
        attempts += 1
        # cheap repair heuristic: drop dangling edges; if still invalid (cycle /
        # too many nodes), give up and let the chain fallback handle it.
        nodeset = set(nodes)
        cur_edges = [(a, b) for (a, b) in cur_edges if a in nodeset and b in nodeset]
        ok, why = is_valid_dag(nodes, cur_edges)
    if ok:
        return list(nodes), cur_edges, "dag", attempts
    # Fallback: sequential chain over (capped) nodes.
    capped = list(nodes)[:N_MAX]
    chain_edges = [(capped[i], capped[i + 1]) for i in range(len(capped) - 1)]
    return capped, chain_edges, "chain-fallback", attempts


def schedule_execute(nodes, edges, u_hat, tau_t):
    """Stages 2 & 3: in-degree-0 frontier routing + topological aggregation.
    Returns (topo_order, routing, parallel_batches)."""
    indeg = {n: 0 for n in nodes}
    adj = {n: [] for n in nodes}
    for a, b in edges:
        indeg[b] += 1
        adj[a].append(b)
    frontier = deque(sorted([n for n in nodes if indeg[n] == 0]))
    topo, routing, batches = [], {}, []
    # batches capture the sets of independent subtasks ready at the same time
    # (i.e. what the DAG scheduler can run in parallel).
    ready = [n for n in nodes if indeg[n] == 0]
    while frontier:
        batch = list(frontier)
        batches.append(sorted(batch))
        next_frontier = deque()
        for _ in range(len(frontier)):
            x = frontier.popleft()
            topo.append(x)
            routing[x] = "cloud" if u_hat[x] > tau_t else "edge"
            for y in adj[x]:
                indeg[y] -= 1
                if indeg[y] == 0:
                    next_frontier.append(y)
        frontier = deque(sorted(next_frontier))
    assert len(topo) == len(nodes), "all nodes must be scheduled"
    return topo, routing, batches


# ---------------------------------------------------------------------------
# Test 1: a valid DAG with independent (parallelizable) subtasks.
#   t1 -> t3, t2 -> t3, t3 -> t5, t4 -> t5   (t1,t2,t4 independent at the root)
# ---------------------------------------------------------------------------
nodes1 = ["t1", "t2", "t3", "t4", "t5"]
edges1 = [("t1", "t3"), ("t2", "t3"), ("t3", "t5"), ("t4", "t5")]
u1 = {"t1": 0.8, "t2": 0.3, "t3": 0.9, "t4": 0.1, "t5": 0.7}
TAU = 0.5
g_nodes, g_edges, mode1, att1 = validate_and_repair(nodes1, edges1)
topo1, route1, batches1 = schedule_execute(g_nodes, g_edges, u1, TAU)
# topological correctness: every edge respected in the order
pos = {n: i for i, n in enumerate(topo1)}
assert all(pos[a] < pos[b] for a, b in g_edges), "topological order violated"
parallel1 = [b for b in batches1 if len(b) > 1]
print("Test 1 (valid DAG):")
print(f"  mode={mode1}  topo order={topo1}")
print(f"  parallel batches (independent subtasks run together)={batches1}")
print(f"  routing={route1}")
assert mode1 == "dag"
assert any(len(b) > 1 for b in batches1), "expected a parallel batch"
print("  -> correct topological order + parallel frontier -> PASS")

# ---------------------------------------------------------------------------
# Test 2: an INVALID DAG (cycle) -> ValidateAndRepair must fall back to a chain.
# ---------------------------------------------------------------------------
nodes2 = ["a", "b", "c"]
edges2 = [("a", "b"), ("b", "c"), ("c", "a")]  # cycle
g2_nodes, g2_edges, mode2, att2 = validate_and_repair(nodes2, edges2)
topo2, route2, batches2 = schedule_execute(
    g2_nodes, g2_edges, {n: 0.6 for n in nodes2}, TAU
)
print("\nTest 2 (cyclic -> invalid):")
print(f"  mode={mode2}  repair_attempts={att2}  chain edges={g2_edges}  topo={topo2}")
assert mode2 == "chain-fallback", "cyclic DAG must fall back to chain"
assert att2 == R_MAX
print("  -> fell back to linear chain after R_max repairs -> PASS")

# ---------------------------------------------------------------------------
# Test 3: too many nodes (> n_max=7) -> invalid -> chain fallback, capped to n_max.
# ---------------------------------------------------------------------------
nodes3 = [f"n{i}" for i in range(9)]
edges3 = [(f"n{i}", f"n{i+1}") for i in range(8)]
g3_nodes, g3_edges, mode3, att3 = validate_and_repair(nodes3, edges3)
print("\nTest 3 (9 nodes > n_max=7):")
print(f"  mode={mode3}  nodes kept={len(g3_nodes)} (capped to n_max={N_MAX})")
assert mode3 == "chain-fallback" and len(g3_nodes) == N_MAX
print("  -> capped to n_max and chained -> PASS")

# ---------------------------------------------------------------------------
# Figure: the Test-1 DAG with edge/cloud routing annotated.
# ---------------------------------------------------------------------------
# simple layered layout by topological batch
layer_of = {}
for li, batch in enumerate(batches1):
    for n in batch:
        layer_of[n] = li
pos_xy = {}
from collections import defaultdict

by_layer = defaultdict(list)
for n, l in layer_of.items():
    by_layer[l].append(n)
for l, ns in by_layer.items():
    for i, n in enumerate(sorted(ns)):
        pos_xy[n] = (l, -(i - (len(ns) - 1) / 2))
fig, ax = plt.subplots(figsize=(7.5, 4.2))
for a, b in g_edges:
    x1, y1 = pos_xy[a]
    x2, y2 = pos_xy[b]
    ax.annotate(
        "",
        xy=(x2, y2),
        xytext=(x1, y1),
        arrowprops=dict(arrowstyle="->", color="#999", lw=1.5),
    )
for n, (x, y) in pos_xy.items():
    col = "#d1495b" if route1[n] == "cloud" else "#2e86ab"
    ax.scatter([x], [y], s=1700, color=col, zorder=3, edgecolors="k")
    ax.text(
        x,
        y,
        f"{n}\n{route1[n]}",
        ha="center",
        va="center",
        color="white",
        fontsize=8,
        zorder=4,
        fontweight="bold",
    )
ax.set_title("Claim 3: Algorithm 1 DAG - cloud (red, u>tau) vs edge (blue) routing")
ax.axis("off")
import matplotlib.patches as mpatches

ax.legend(
    handles=[
        mpatches.Patch(color="#d1495b", label="cloud (u_hat > tau_t)"),
        mpatches.Patch(color="#2e86ab", label="edge"),
    ],
    loc="lower right",
)
plt.tight_layout()
fig.savefig(os.path.join(FIG, "claim3_dag.png"), dpi=130)
plt.close(fig)

result = {
    "claim": "3 - dependency-aware DAG decomposition + 3-stage pipeline (Algorithm 1)",
    "config": {"n_max": N_MAX, "R_max": R_MAX},
    "note": "Single Algorithm 1; Decompose/ValidateAndRepair = Def C.2 + XML planner "
    "prompt (Fig 6), not a numbered Algorithm 2. Scheduling mechanism reproduced "
    "on CPU; LLM planner quality out of scope (needs an LLM).",
    "test1_valid_dag": {
        "nodes": nodes1,
        "edges": edges1,
        "mode": mode1,
        "topo_order": topo1,
        "parallel_batches": batches1,
        "routing": route1,
        "has_parallel_frontier": any(len(b) > 1 for b in batches1),
    },
    "test2_cycle_fallback": {
        "mode": mode2,
        "repair_attempts": att2,
        "chain_edges": g2_edges,
    },
    "test3_too_many_nodes": {
        "input_nodes": len(nodes3),
        "mode": mode3,
        "nodes_kept": len(g3_nodes),
        "n_max": N_MAX,
    },
    "figure": "figs/claim3_dag.png",
    "verdict": "PASS: valid DAG scheduled in correct topological order with a parallel "
    "in-degree-0 frontier and per-subtask edge/cloud routing; invalid DAGs "
    "(cycle, >n_max) fall back to a linear chain after R_max repairs.",
}
with open(os.path.join(OUT, "claim3.json"), "w") as f:
    json.dump(result, f, indent=2)
print("\nVERDICT: PASS")