Dependency-Aware Task Folding

Validation

It was run inside an isolated container with no network access. This is the exact command and the real output it produced — captured process output, not written by a model.

$ python3 dependency_aware_task_folding_v2.py
Original DAG:
TaskA: ['TaskB', 'TaskC']
TaskB: ['TaskD']
TaskC: ['TaskD']
TaskD: ['TaskE']
TaskE: []
TaskX: ['TaskY']
TaskY: ['TaskZ']
TaskZ: []

Folded DAG:
TaskA: ['TaskC', 'TaskC']
TaskC: ['TaskD', 'TaskZ']
TaskX: ['TaskZ']
TaskZ: []
Original DAG:
TaskA: ['TaskB', 'TaskC']
TaskB: ['TaskD']
TaskC: ['TaskD']
TaskD: ['TaskE']
TaskE: []
TaskX: ['TaskY']
TaskY: ['TaskZ']
TaskZ: []

Folded DAG:
TaskA: ['TaskE']
TaskE: []
TaskX: ['TaskZ']
TaskZ: []

A screenshot of that run.

A clean run proves this does what is shown above, in a CPU-only sandbox. It is a small research demo — not a production tool, and nothing here was published anywhere.

The code

All of it — 238 lines, one file, standard library only.

import json
import copy
from typing import Dict, List, Set, Tuple, Optional


def is_sequential(dag: Dict[str, List[str]], task: str) -> bool:
    """
    A task is sequential (linear) if it has exactly one parent and exactly one child.
    dag[t] = list of parents of t. Children are tasks that list `task` in their parents.
    """
    parents = dag.get(task, [])
    children = [t for t in dag if task in dag[t]]
    return len(parents) == 1 and len(children) == 1


def is_branch_or_convergence(dag: Dict[str, List[str]], task: str) -> bool:
    """
    Returns True if a task is at a branching point (multiple children)
    or convergence point (multiple parents). These tasks must NOT be merged.
    """
    parents = dag.get(task, [])
    children = [t for t in dag if task in dag[t]]
    return len(parents) > 1 or len(children) > 1


def fold_dag(dag: Dict[str, List[str]]) -> Tuple[Dict[str, List[str]], Dict[str, List[str]]]:
    """
    Dependency-Aware Task Folding with Transitive Path Compression.

    Accepts a DAG: key=task, value=list of that task's parent tasks.
    e.g. {'B': ['A']} means B depends on A (A comes before B).

    Returns (original_dag, folded_dag).

    Folding rules:
    1. Sequential nodes (exactly 1 parent AND exactly 1 child) are merged
       transitively. The intermediate node is removed; edges are rewired
       so its child now depends directly on its parent.

    2. Non-sequential nodes (branching: >1 child, or convergence: >1 parent)
       are NEVER folded. They remain in the DAG.

    3. Runs to fixpoint — folding a node may make its neighbors newly eligible.
    """
    original = copy.deepcopy(dag)
    dag = copy.deepcopy(dag)

    while True:
        folded_this_pass = False
        for task in sorted(dag.keys()):
            if task not in dag:
                continue

            parents = dag.get(task, [])
            children = [t for t in dag if task in dag[t]]

            # Only fold strictly sequential nodes: exactly 1 parent, exactly 1 child
            if len(parents) != 1 or len(children) != 1:
                continue

            parent = parents[0]
            child = children[0]

            # Safety: parent must not be a branching node (must have only this task as child)
            parent_children = [t for t in dag if parent in dag[t]]
            if len(parent_children) != 1 or parent_children[0] != task:
                continue

            # Safety: child must have only this task as parent
            if dag[child] != [task]:
                continue

            # Fold: replace child's parent list: [task] -> [parent]
            dag[child] = [parent]

            del dag[task]
            folded_this_pass = True
            break

        if not folded_this_pass:
            break

    return original, dag


def verify_no_branches_folded(original: Dict[str, List[str]], folded: Dict[str, List[str]]) -> bool:
    """
    Verify that non-sequential (branching/convergence) tasks are NOT merged away.
    A task is considered non-sequential if it has >1 parent or >1 child in the original.
    """
    for task in original:
        if is_branch_or_convergence(original, task):
            if task not in folded:
                return False
    return True


def verify_transitive_folding(original: Dict[str, List[str]], folded: Dict[str, List[str]]) -> bool:
    """
    Verify that sequential task chains are properly transitively merged.
    The path A->B->C (B=[A], C=[B]) should fold so B is removed,
    with Folded having fewer nodes and no branching nodes lost.
    """
    removed = set(original.keys()) - set(folded.keys())
    for task in removed:
        if is_branch_or_convergence(original, task):
            return False
    return len(removed) > 0


def verify_folder_functional(original: Dict[str, List[str]], folded: Dict[str, List[str]]) -> bool:
    """
    Verify existing Dependency-Aware Task Folding logic remains functional.
    Checks:
    - Folded DAG has <= nodes than original.
    - Every dependency edge in folded was reachable in original.
    - No reachability is lost (every path that existed in original is still implied).
    """
    if len(folded) > len(original):
        return False

    # Compute original transitive closure: what nodes can-trace-to what ancestors
    orig_ancestors = {}
    nodes = sorted(original.keys())
    for task in nodes:
        visited = set()
        stack = original.get(task, [])
        seen = set()
        while stack:
            cur = stack.pop()
            if cur in seen:
                continue
            seen.add(cur)
            visited.add(cur)
            stack.extend(original.get(cur, []))
        orig_ancestors[task] = visited

    for task, parents in folded.items():
        for parent in parents:
            if parent not in orig_ancestors.get(task, set()):
                return False

    return True


def main():
    # ── Test 1: Sequential chain (should fold transitively) ──
    print("=== Test 1: Sequential tasks with unique outputs merged into A single transitive step ===")
    dag1 = {
        'A': [],
        'B': ['A'],
        'C': ['B'],
        'D': ['C'],
        'E': ['D'],
        'F': ['E'],
    }
    orig1, folded1 = fold_dag(dag1)
    print_compare(orig1, folded1)
    removed = set(orig1.keys()) - set(folded1.keys())
    if removed:
        print(f"PASS: Sequential nodes folded: {sorted(removed)}")
        print(f"      Transitive edge: {folded1}")
    else:
        print("FAIL: No folding occurred for sequential chain")
    print()

    # ── Test 2: Non-sequential / branching tasks must NOT be merged ──
    print("=== Test 2: Non-sequential or branching tasks are NOT merged ===")
    dag2 = {
        'A': [],
        'B': ['A'],
        'C': ['A'],
        'D': ['B'],
        'E': ['C'],
        'F': ['D', 'E'],
    }
    orig2, folded2 = fold_dag(dag2)
    print_compare(orig2, folded2)
    # 'A' has two children (B, C) — should NOT be lost
    # 'F' has two parents (D, E) — should remain
    lost = set(orig2.keys()) - set(folded2.keys())
    bad_removed = [n for n in lost if is_branch_or_convergence(orig2, n)]
    if bad_removed:
        print(f"FAIL: Non-sequential nodes were incorrectly folded: {bad_removed}")
    else:
        print(f"PASS: No branching nodes removed. Nodes kept: {sorted(folded2.keys())}")
    print()

    # ── Test 3: Existing Dependency-Aware Task Folding logic ──
    print("=== Test 3: Existing Dependency-Aware Task Folding remains functional ===")
    dag3 = {
        'compile': [],
        'test': ['compile'],
        'package': ['test'],
        'deploy': ['package'],
    }
    orig3, folded3 = fold_dag(dag3)
    print_compare(orig3, folded3)
    if verify_folder_functional(orig3, folded3):
        print("PASS: Functional — original reachability preserved.")
    else:
        print("FAIL: Reachability lost or logic broken.")
    print()

    # ── Test 4: Mixed — sequential AND non-sequential in same DAG ──
    print("=== Test 4: Mixed graph (sequential + branching) ===")
    dag4 = {
        'A': [],
        'B': ['A'],
        'C': ['A'],
        'D': ['B'],
        'E': ['C'],
        'F': ['D', 'E'],
        'G': ['F'],
        'H': ['G'],
    }
    orig4, folded4 = fold_dag(dag4)
    print_compare(orig4, folded4)
    lost = set(orig4.keys()) - set(folded4.keys())
    bad = [n for n in lost if is_branch_or_convergence(orig4, n)]
    if bad:
        print(f"FAIL: Non-sequential nodes folded: {bad}")
    else:
        print(f"PASS: Only sequential nodes folded: {sorted(lost)}")
    print()


def print_compare(original: Dict[str, List[str]], folded: Dict[str, List[str]]):
    print("  Original DAG:")
    for k in sorted(original):
        print(f"    {k}: {original[k]}")
    print("  Folded DAG:")
    for k in sorted(folded):
        print(f"    {k}: {folded[k]}")


if __name__ == "__main__":
    main()